# process inputs

sequence = str(sequence)

sequence = re.sub("[^A-Z:/]", "", sequence.upper())

sequence = re.sub(":+",":",sequence)

sequence = re.sub("/+","/",sequence)

sequence = re.sub("^[:/]+","",sequence)

sequence = re.sub("[:/]+$","",sequence)

jobname = re.sub(r'\W+', '', jobname)

homooligomer = str(homooligomer)

homooligomer = re.sub("[:/]+",":",homooligomer)

homooligomer = re.sub("^[:/]+","",homooligomer)

homooligomer = re.sub("[:/]+$","",homooligomer)

if len(homooligomer) == 0: homooligomer = "1"

homooligomer = re.sub("[^0-9:]", "", homooligomer)

# define inputs

I = {"ori_sequence":sequence,

"sequence":sequence.replace("/","").replace(":",""),

"seqs":sequence.replace("/","").split(":"),

"homooligomer":homooligomer,

"homooligomers":[int(h) for h in homooligomer.split(":")],

"msas":[], "deletion_matrices":[]}

# adjust homooligomer option

if len(I["seqs"]) != len(I["homooligomers"]):

if len(I["homooligomers"]) == 1:

I["homooligomers"] = [I["homooligomers"][0]] * len(I["seqs"])

else:

if verbose:

print("WARNING: Mismatch between number of breaks ':' in 'sequence' and 'homooligomer' definition")

while len(I["seqs"]) > len(I["homooligomers"]):

I["homooligomers"].append(1)

I["homooligomers"] = I["homooligomers"][:len(I["seqs"])]

I["homooligomer"] = ":".join([str(h) for h in I["homooligomers"]])

# define full sequence being modelled

I["full_sequence"] = ''.join([s*h for s,h in zip(I["seqs"],I["homooligomers"])])

I["lengths"] = [len(seq) for seq in I["seqs"]]

# prediction directory

if output_dir is None:

I["output_dir"] = 'prediction_' + jobname + '_' + cf.get_hash(I["full_sequence"])[:5]

else:

I["output_dir"] = output_dir

os.makedirs(I["output_dir"], exist_ok=True)

# delete existing files in working directory

if clean:

for f in os.listdir(I["output_dir"]):

os.remove(os.path.join(I["output_dir"], f))

if verbose and len(I["full_sequence"]) > 1400:

print(f"WARNING: For a typical Google-Colab-GPU (16G) session, the max total length is ~1400 residues. You are at {len(I['full_sequence'])}!")

print(f"Run Alphafold may crash, unless you trim to the protein(s) to a short length. (See trim options below).")

if verbose:

print(f"homooligomer: {I['homooligomer']}")

print(f"total_length: {len(I['full_sequence'])}")

print(f"output_dir: {I['output_dir']}")

fasta_path = f"{prefix}.fasta"

with open(fasta_path, 'wt') as f:

f.write(f'>query\n{sequence}')

pickled_msa_path = f"{prefix}.jackhmmer.pickle"

if os.path.isfile(pickled_msa_path):

msas_dict = pickle.load(open(pickled_msa_path,"rb"))

msas, deletion_matrices, names = (msas_dict[k] for k in ['msas', 'deletion_matrices', 'names'])

full_msa = []

for msa in msas:

full_msa += msa

else:

# --- Find the closest source ---

test_url_pattern = 'https://storage.googleapis.com/alphafold-colab{:s}/latest/uniref90_2021_03.fasta.1'

ex = futures.ThreadPoolExecutor(3)

def fetch(source):

request.urlretrieve(test_url_pattern.format(source))

return source

fs = [ex.submit(fetch, source) for source in ['', '-europe', '-asia']]

source = None

for f in futures.as_completed(fs):

source = f.result()

ex.shutdown()

break

dbs = []

num_jackhmmer_chunks = {'uniref90': 59, 'smallbfd': 17, 'mgnify': 71}

total_jackhmmer_chunks = sum(num_jackhmmer_chunks.values())

disable_tqdm = not verbose

with tqdm.notebook.tqdm(total=total_jackhmmer_chunks, bar_format=TQDM_BAR_FORMAT, disable=disable_tqdm) as pbar:

def jackhmmer_chunk_callback(i):

pbar.update(n=1)

pbar.set_description('Searching uniref90')

jackhmmer_uniref90_runner = jackhmmer.Jackhmmer(

binary_path=jackhmmer_binary_path,

database_path=f'https://storage.googleapis.com/alphafold-colab{source}/latest/uniref90_2021_03.fasta',

get_tblout=True,

num_streamed_chunks=num_jackhmmer_chunks['uniref90'],

streaming_callback=jackhmmer_chunk_callback,

z_value=135301051)

dbs.append(('uniref90', jackhmmer_uniref90_runner.query(fasta_path)))

pbar.set_description('Searching smallbfd')

jackhmmer_smallbfd_runner = jackhmmer.Jackhmmer(

binary_path=jackhmmer_binary_path,

database_path=f'https://storage.googleapis.com/alphafold-colab{source}/latest/bfd-first_non_consensus_sequences.fasta',

get_tblout=True,

num_streamed_chunks=num_jackhmmer_chunks['smallbfd'],

streaming_callback=jackhmmer_chunk_callback,

z_value=65984053)

dbs.append(('smallbfd', jackhmmer_smallbfd_runner.query(fasta_path)))

pbar.set_description('Searching mgnify')

jackhmmer_mgnify_runner = jackhmmer.Jackhmmer(

binary_path=jackhmmer_binary_path,

database_path=f'https://storage.googleapis.com/alphafold-colab{source}/latest/mgy_clusters_2019_05.fasta',

get_tblout=True,

num_streamed_chunks=num_jackhmmer_chunks['mgnify'],

streaming_callback=jackhmmer_chunk_callback,

z_value=304820129)

dbs.append(('mgnify', jackhmmer_mgnify_runner.query(fasta_path)))

# --- Extract the MSAs and visualize ---

# Extract the MSAs from the Stockholm files.

# NB: deduplication happens later in pipeline.make_msa_features.

mgnify_max_hits = 501

msas = []

deletion_matrices = []

names = []

for db_name, db_results in dbs:

unsorted_results = []

for i, result in enumerate(db_results):

msa, deletion_matrix, target_names = parsers.parse_stockholm(result['sto'])

e_values_dict = parsers.parse_e_values_from_tblout(result['tbl'])

e_values = [e_values_dict[t.split('/')[0]] for t in target_names]

zipped_results = zip(msa, deletion_matrix, target_names, e_values)

if i != 0:

# Only take query from the first chunk

zipped_results = [x for x in zipped_results if x[2] != 'query']

unsorted_results.extend(zipped_results)

sorted_by_evalue = sorted(unsorted_results, key=lambda x: x[3])

db_msas, db_deletion_matrices, db_names, _ = zip(*sorted_by_evalue)

if db_msas:

if db_name == 'mgnify':

db_msas = db_msas[:mgnify_max_hits]

db_deletion_matrices = db_deletion_matrices[:mgnify_max_hits]

db_names = db_names[:mgnify_max_hits]

msas.append(db_msas)

deletion_matrices.append(db_deletion_matrices)

names.append(db_names)

msa_size = len(set(db_msas))

print(f'{msa_size} Sequences Found in {db_name}')

pickle.dump({"msas":msas,

"deletion_matrices":deletion_matrices,

"names":names}, open(pickled_msa_path,"wb"))

pair_mode="unpaired", pair_cov=50, pair_qid=20,

hhfilter_loc="hhfilter", reformat_loc="reformat.pl", TMP_DIR="tmp",

custom_msa=None, precomputed=None,

mmseqs_host_url="https://a3m.mmseqs.com",

verbose=True):

# make temp directory

os.makedirs(TMP_DIR, exist_ok=True)

# clear previous inputs

I["msas"] = []

I["deletion_matrices"] = []

if add_custom_msa:

if IN_COLAB:

print(f"upload custom msa in '{msa_format}' format")

msa_dict = files.upload()

lines = msa_dict[list(msa_dict.keys())[0]].decode()

input_file = os.path.join(I["output_dir"],f"upload.{msa_format}")

with open(input_file,"w") as tmp_upload:

tmp_upload.write(lines)

else:

input_file = custom_msa

if input_file is None or not os.path.isfile(input_file):

raise ValueError("ERROR: `custom_msa` undefined")

else:

# convert to a3m

output_file = os.path.join(I["output_dir"],f"upload.a3m")

os.system(f"{reformat_loc} {msa_format} a3m {input_file} {output_file}")

# parse

msa, mtx = parsers.parse_a3m(open(output_file,"r").read())

I["msas"].append(msa)

I["deletion_matrices"].append(mtx)

if len(I["msas"][0][0]) != len(I["sequence"]):

raise ValueError("ERROR: the length of msa does not match input sequence")

if msa_method == "precomputed":

if IN_COLAB:

print("upload precomputed pickled msa from previous run")

uploaded_dict = files.upload()

uploaded_filename = list(uploaded_dict.keys())[0]

I.update(pickle.loads(uploaded_dict[uploaded_filename]))

elif precomputed is None:

raise ValueError("ERROR: `precomputed` undefined")

else:

I.update(pickle.load(open(precomputed,"rb")))

elif msa_method == "single_sequence":

if len(I["msas"]) == 0:

I["msas"].append([I["sequence"]])

I["deletion_matrices"].append([[0]*len(I["sequence"])])

else:

_blank_seq = ["-" * L for L in I["lengths"]]

_blank_mtx = [[0] * L for L in I["lengths"]]

def _pad(ns,vals,mode):

if mode == "seq": _blank = _blank_seq.copy()

if mode == "mtx": _blank = _blank_mtx.copy()

if isinstance(ns, list):

for n,val in zip(ns,vals): _blank[n] = val

else: _blank[ns] = vals

if mode == "seq": return "".join(_blank)

if mode == "mtx": return sum(_blank,[])

if len(I["seqs"]) == 1 or "unpaired" in pair_mode:

# gather msas

if msa_method == "mmseqs2":

prefix = cf.get_hash(I["sequence"])

prefix = os.path.join(TMP_DIR,prefix)

print(f"running mmseqs2")

A3M_LINES = cf.run_mmseqs2(I["seqs"], prefix, use_filter=True, host_url=mmseqs_host_url)

for n, seq in enumerate(I["seqs"]):

# tmp directory

prefix = cf.get_hash(seq)

prefix = os.path.join(TMP_DIR,prefix)

if msa_method == "mmseqs2":

# run mmseqs2

a3m_lines = A3M_LINES[n]

msa, mtx = parsers.parse_a3m(a3m_lines)

msas_, mtxs_ = [msa],[mtx]

elif msa_method == "jackhmmer":

print(f"running jackhmmer on seq_{n}")

# run jackhmmer

msas_, mtxs_, names_ = ([sum(x,())] for x in run_jackhmmer(seq, prefix))

# pad sequences

for msa_,mtx_ in zip(msas_,mtxs_):

msa,mtx = [I["sequence"]],[[0]*len(I["sequence"])]

for s,m in zip(msa_,mtx_):

msa.append(_pad(n,s,"seq"))

mtx.append(_pad(n,m,"mtx"))

I["msas"].append(msa)

I["deletion_matrices"].append(mtx)

# PAIR_MSA

if len(I["seqs"]) > 1 and (pair_mode == "paired" or pair_mode == "unpaired+paired"):

print("attempting to pair some sequences...")

if msa_method == "mmseqs2":

prefix = cf.get_hash(I["sequence"])

prefix = os.path.join(TMP_DIR,prefix)

print(f"running mmseqs2_noenv_nofilter on all seqs")

A3M_LINES = cf.run_mmseqs2(I["seqs"], prefix, use_env=False, use_filter=False, host_url=mmseqs_host_url)

_data = []

for a in range(len(I["seqs"])):

print(f"prepping seq_{a}")

_seq = I["seqs"][a]

_prefix = os.path.join(TMP_DIR,cf.get_hash(_seq))

if msa_method == "mmseqs2":

a3m_lines = A3M_LINES[a]

_msa, _mtx, _lab = pairmsa.parse_a3m(a3m_lines,

filter_qid=pair_qid/100,

filter_cov=pair_cov/100)

elif msa_method == "jackhmmer":

_msas, _mtxs, _names = run_jackhmmer(_seq, _prefix)

_msa, _mtx, _lab = pairmsa.get_uni_jackhmmer(_msas[0], _mtxs[0], _names[0],

filter_qid=pair_qid/100,

filter_cov=pair_cov/100)

if len(_msa) > 1:

_data.append(pairmsa.hash_it(_msa, _lab, _mtx, call_uniprot=False))

else:

_data.append(None)

Ln = len(I["seqs"])

O = [[None for _ in I["seqs"]] for _ in I["seqs"]]

for a in range(Ln):

if _data[a] is not None:

for b in range(a+1,Ln):

if _data[b] is not None:

print(f"attempting pairwise stitch for {a} {b}")

O[a][b] = pairmsa._stitch(_data[a],_data[b])

_seq_a, _seq_b, _mtx_a, _mtx_b = (*O[a][b]["seq"],*O[a][b]["mtx"])

# filter to remove redundant sequences

ok = []

with open(f"{TMP_DIR}/tmp.fas","w") as fas_file:

fas_file.writelines([f">{n}\n{a+b}\n" for n,(a,b) in enumerate(zip(_seq_a,_seq_b))])

os.system(f"{hhfilter_loc} -maxseq 1000000 -i {TMP_DIR}/tmp.fas -o {TMP_DIR}/tmp.id90.fas -id 90")

for line in open(f"{TMP_DIR}/tmp.id90.fas","r"):

if line.startswith(">"): ok.append(int(line[1:]))

if verbose:

print(f"found {len(_seq_a)} pairs ({len(ok)} after filtering)")

if len(_seq_a) > 0:

msa,mtx = [I["sequence"]],[[0]*len(I["sequence"])]

for s_a,s_b,m_a,m_b in zip(_seq_a, _seq_b, _mtx_a, _mtx_b):

msa.append(_pad([a,b],[s_a,s_b],"seq"))

mtx.append(_pad([a,b],[m_a,m_b],"mtx"))

I["msas"].append(msa)

I["deletion_matrices"].append(mtx)

# save MSA as pickle

pickle.dump({"msas":I["msas"],"deletion_matrices":I["deletion_matrices"]},

open(os.path.join(I["output_dir"],"msa.pickle"),"wb"))

'''

if ori_seq is None: ori_seq = msas[0][0]

seqs = ori_seq.replace("/","").split(":")

L_ini = 0

chain_idx = {}

idx_chain = []

for chain,seq in zip(ascii_uppercase,seqs):

L = len(seq)

chain_idx[chain] = dict(zip(range(L),range(L_ini,L_ini+L)))

idx_chain += [f"{chain}{i+1}" for i in range(L)]

L_ini += L

global_idx = dict(zip(range(L_ini),range(L_ini)))

mode = "keeping" if inverse else "trimming"

trim_set = []

for idx in trim.split(","):

i,j = idx.split("-") if "-" in idx else (idx,"")

# set index reference frame

trim_idx_i = trim_idx_j = global_idx

if i != "" and i[0] in ascii_uppercase:

trim_idx_i,i = chain_idx[i[0]], i[1:]

if j != "" and j[0] in ascii_uppercase:

trim_idx_j,j = chain_idx[j[0]], j[1:]

# set which positions to trim

if "-" in idx:

i = trim_idx_i[int(i)-1] if i != "" else trim_idx_i[0]

j = trim_idx_j[int(j)-1] if j != "" else trim_idx_j[len(trim_idx_j) - 1]

trim_set += list(range(i,j+1))

print(f"{mode} positions: {idx_chain[i]}-{idx_chain[j]}")

else:

i = trim_idx_i[int(i)-1]

trim_set.append(i)

print(f"{mode} position: {idx_chain[i]}")

# deduplicate list

trim_set = set(trim_set)

if inverse:

trim_set = set(range(L_ini)) ^ trim_set

trim_set = sorted(list(trim_set))

# trim MSA

mod_msas, mod_mtxs = [],[]

for msa, mtx in zip(msas, deletion_matrices):

mod_msa = np.delete([list(s) for s in msa], trim_set, 1)

ok = (mod_msa != "-").sum(-1) > 0

mod_msas.append(["".join(s) for s in mod_msa[ok]])

mod_mtx = np.asarray(mtx)[ok]

mod_mtxs.append(np.delete(mod_mtx, trim_set, 1).tolist())

# trim original sequence

mod_idx = []

mod_chain = []

mod_ori_seq = []

for n,a in enumerate(ori_seq.replace("/","").replace(":","")):

if n not in trim_set:

mod_ori_seq.append(a)

mod_idx.append(n)

mod_chain.append(idx_chain[n][0])

if len(mod_idx) > 1:

if mod_chain[-1] != mod_chain[-2]:

mod_ori_seq[-1] = ":"

mod_ori_seq.append(a)

elif (mod_idx[-1] - mod_idx[-2]) > 1:

mod_ori_seq[-1] = "/"

mod_ori_seq.append(a)

mod_ori_seq = "".join(mod_ori_seq)

chains = sorted([ascii_uppercase.index(a) for a in set(mod_chain)])

return {"msas":mod_msas, "deletion_matrices":mod_mtxs,

if ori_seq is None: ori_seq = msas[0][0]

seqs = ori_seq.replace("/","").split(":")

ref_seq_ = np.array(list("".join(seqs)))

new_msas,new_mtxs = [],[]

L = np.asarray([len(seq) for seq in seqs])

Ln = np.cumsum(np.append(0,L))

for msa, mtx in zip(msas, deletion_matrices):

msa_ = np.asarray([list(seq) for seq in msa])

# coverage (non-gap characters)

cov_ = msa_ != "-"

# sequence identity to query

qid_ = msa_ == ref_seq_

# split by protein (for protein complexes)

cov__ = np.stack([cov_[:,Ln[i]:Ln[i+1]].sum(-1) for i in range(len(seqs))],-1)

qid__ = np.stack([qid_[:,Ln[i]:Ln[i+1]].sum(-1) for i in range(len(seqs))],-1)

not_empty__ = cov__ > 0

ok = []

for n in range(len(msa)):

m = not_empty__[n]

if m.sum() > 0:

q = qid__[n][m].sum() / cov__[n][m].sum()

c = cov__[n][m].sum() / L[m].sum()

if q > qid and c > cov:

ok.append(n)

new_msas.append([msa[n] for n in ok])

new_mtxs.append([mtx[n] for n in ok])

trim = re.sub("[^0-9A-Z,-]", "", trim.upper())

trim = re.sub(",+",",",trim)

trim = re.sub("^[,]+","",trim)

trim = re.sub("[,]+$","",trim)

if trim != "" or cov > 0 or qid > 0:

mod_I = dict(I)

if trim != "":

mod_I.update(trim_inputs(trim, mod_I["msas"], mod_I["deletion_matrices"],

mod_I["ori_sequence"], inverse=trim_inverse))

mod_I["homooligomers"] = [mod_I["homooligomers"][c] for c in mod_I["chains"]]

mod_I["sequence"] = mod_I["ori_sequence"].replace("/","").replace(":","")

mod_I["seqs"] = mod_I["ori_sequence"].replace("/","").split(":")

mod_I["full_sequence"] = "".join([s*h for s,h in zip(mod_I["seqs"], mod_I["homooligomers"])])

new_length = len(mod_I["full_sequence"])

if verbose:

print(f"total_length: '{new_length}' after trimming")

if cov > 0 or qid > 0:

mod_I.update(cov_qid_filter(mod_I["msas"], mod_I["deletion_matrices"],

mod_I["ori_sequence"], cov=cov/100, qid=qid/100))

return mod_I

else:

def _placeholder_template_feats(num_templates_, num_res_):

return {

'template_aatype': np.zeros([num_templates_, num_res_, 22], np.float32),

'template_all_atom_masks': np.zeros([num_templates_, num_res_, 37], np.float32),

'template_all_atom_positions': np.zeros([num_templates_, num_res_, 37, 3], np.float32),

'template_domain_names': np.zeros([num_templates_], np.float32),

'template_sum_probs': np.zeros([num_templates_], np.float32),

}

# delete old files

if clean:

for f in os.listdir(I["output_dir"]):

if "rank_" in f: os.remove(os.path.join(I["output_dir"], f))

if len(I["msas"]) == 0:

print("WARNING: no MSA found, switching to 'single_sequence' mode")

I["msas"].append([I["sequence"]])

I["deletion_matrices"].append([[0]*len(I["sequence"])])

# homooligomerize

lengths = [len(seq) for seq in I["seqs"]]

msas_mod, deletion_matrices_mod = cf.homooligomerize_heterooligomer(I["msas"], I["deletion_matrices"],

lengths, I["homooligomers"])

# define input features

num_res = len(I["full_sequence"])

feature_dict = {}

feature_dict.update(pipeline.make_sequence_features(I["full_sequence"], 'test', num_res))

feature_dict.update(pipeline.make_msa_features(msas_mod, deletion_matrices=deletion_matrices_mod))

feature_dict.update(_placeholder_template_feats(0, num_res))

# set chainbreaks

Ls = []

for seq,h in zip(I["ori_sequence"].split(":"), I["homooligomers"]):

Ls += [len(s) for s in seq.split("/")] * h

Ls_plot = []

for seq,h in zip(I["seqs"], I["homooligomers"]):

Ls_plot += [len(seq)] * h

feature_dict['residue_index'] = cf.chain_break(feature_dict['residue_index'], Ls)

feature_dict['Ls'] = Ls_plot

feature_dict['output_dir'] = I["output_dir"]

'''pad input features'''

opt = runner["opt"]

cfg = runner["model"].config

shape_schema = {k:[None]+v for k,v in dict(cfg.data.eval.feat).items()}

pad_size_map = {

shape_placeholders.NUM_RES: opt["L"],

shape_placeholders.NUM_MSA_SEQ: cfg.data.eval.max_msa_clusters,

shape_placeholders.NUM_EXTRA_SEQ: cfg.data.common.max_extra_msa,

shape_placeholders.NUM_TEMPLATES: 0,

}

for k, v in feat.items():

# Don't transfer this to the accelerator.

if k == 'extra_cluster_assignment':

continue

shape = list(v.shape)

schema = shape_schema[k]

assert len(shape) == len(schema), (

f'Rank mismatch between shape and shape schema for {k}: '

f'{shape} vs {schema}')

pad_size = [pad_size_map.get(s2, None) or s1 for (s1, s2) in zip(shape, schema)]

padding = [(0, p - tf.shape(v)[i]) for i, p in enumerate(pad_size)]

if padding:

feat[k] = tf.pad(v, padding, name=f'pad_to_fixed_{k}')

feat[k].set_shape(pad_size)

'''remove all data from device'''

backend = jax.lib.xla_bridge.get_backend(device)

if hasattr(backend,'live_buffers'):

for buf in backend.live_buffers():

"use_ptm":True, "use_turbo":True,

"max_recycles":3, "tol":0, "num_ensemble":1,

"max_msa_clusters":512, "max_extra_msa":1024,

# setup the [opt]ions

if opt is None:

opt = OPT_DEFAULT.copy()

else:

for k in OPT_DEFAULT:

if k not in opt: opt[k] = OPT_DEFAULT[k]

# if old_runner not defined or [opt]ions changed, start new runner

if old_runner is None or old_runner["opt"] != opt:

clear_mem()

name = f"{model_name}_ptm" if opt["use_ptm"] else model_name

cfg = config.model_config(name)

if opt["use_turbo"]:

if opt["N"] is None:

cfg.data.eval.max_msa_clusters = opt["max_msa_clusters"]

cfg.data.common.max_extra_msa = opt["max_extra_msa"]

else:

msa_clusters = min(opt["N"], opt["max_msa_clusters"])

cfg.data.eval.max_msa_clusters = msa_clusters

cfg.data.common.max_extra_msa = max(min(opt["N"] - msa_clusters, opt["max_extra_msa"]),1)

cfg.data.common.num_recycle = opt["max_recycles"]

cfg.model.num_recycle = opt["max_recycles"]

cfg.model.recycle_tol = opt["tol"]

cfg.data.eval.num_ensemble = opt["num_ensemble"]

params = data.get_model_haiku_params(name, params_loc)

return {"model":model.RunModel(cfg, params, is_training=opt["is_training"]), "opt":opt}

else:

pad_feats=False, rank_by="pLDDT", show_images=True, params_loc='./alphafold/data', verbose=True):

def do_subsample_msa(F, random_seed=0):

'''subsample msa to avoid running out of memory'''

N = len(F["msa"])

L = len(F["residue_index"])

N_ = int(3E7/L)

if N > N_:

if verbose:

print(f"whhhaaa... too many sequences ({N}) subsampling to {N_}")

np.random.seed(random_seed)

idx = np.append(0,np.random.permutation(np.arange(1,N)))[:N_]

F_ = {}

F_["msa"] = F["msa"][idx]

F_["deletion_matrix_int"] = F["deletion_matrix_int"][idx]

F_["num_alignments"] = np.full_like(F["num_alignments"],N_)

for k in F.keys():

if k not in F_: F_[k] = F[k]

return F_

else:

return F

def parse_results(prediction_result, processed_feature_dict, r, t, num_res):

'''parse results and convert to numpy arrays'''

to_np = lambda a: np.asarray(a)

def class_to_np(c):

class dict2obj():

def __init__(self, d):

for k,v in d.items(): setattr(self, k, to_np(v))

return dict2obj(c.__dict__)

dist_bins = jax.numpy.append(0,prediction_result["distogram"]["bin_edges"])

dist_logits = prediction_result["distogram"]["logits"][:num_res,:][:,:num_res]

dist_mtx = dist_bins[dist_logits.argmax(-1)]

contact_mtx = jax.nn.softmax(dist_logits)[:,:,dist_bins < 8].sum(-1)

b_factors = prediction_result['plddt'][:,None] * prediction_result['structure_module']['final_atom_mask']

p = protein.from_prediction(processed_feature_dict, prediction_result, b_factors=b_factors)

plddt = prediction_result['plddt'][:num_res]

out = {"unrelaxed_protein": class_to_np(p),

"plddt": to_np(plddt),

"pLDDT": to_np(plddt.mean()),

"dists": to_np(dist_mtx),

"adj": to_np(contact_mtx),

"recycles":to_np(r),

"tol":to_np(t)}

if "ptm" in prediction_result:

out["pae"] = to_np(prediction_result['predicted_aligned_error'][:num_res,:][:,:num_res])

out["pTMscore"] = to_np(prediction_result['ptm'])

return out

num_res = len(feature_dict["residue_index"])

# if [opt]ions not defined

if opt is None:

opt = OPT_DEFAULT.copy()

opt["N"] = len(feature_dict["msa"])

opt["L"] = num_res

else:

for k in OPT_DEFAULT.keys():

if k not in opt: opt[k] = OPT_DEFAULT[k]

model_names = ['model_1', 'model_2', 'model_3', 'model_4', 'model_5'][:num_models]

total = len(model_names) * num_samples

outs = {}

def do_report(key):

o = outs[key]

if verbose:

line = f"{key} recycles:{o['recycles']} tol:{o['tol']:.2f} pLDDT:{o['pLDDT']:.2f}"

if 'pTMscore' in o:

line += f" pTMscore:{o['pTMscore']:.2f}"

print(line)

if show_images:

fig = cf.plot_protein(o['unrelaxed_protein'], Ls=feature_dict["Ls"], dpi=100)

plt.show()

tmp_pdb_path = os.path.join(feature_dict["output_dir"],f'unranked_{key}_unrelaxed.pdb')

pdb_lines = protein.to_pdb(o['unrelaxed_protein'])

with open(tmp_pdb_path, 'w') as f: f.write(pdb_lines)

disable_tqdm = not verbose

with tqdm.notebook.tqdm(total=total, bar_format=TQDM_BAR_FORMAT, disable=disable_tqdm) as pbar:

if opt["use_turbo"]:

if runner is None:

runner = prep_model_runner(opt,params_loc=params_loc)

# go through each random_seed

for seed in range(num_samples):

# prep input features

feat = do_subsample_msa(feature_dict, random_seed=seed) if subsample_msa else feature_dict

processed_feature_dict = runner["model"].process_features(feat, random_seed=seed)

if pad_feats:

processed_feature_dict = make_fixed_size(processed_feature_dict, runner)

# go through each model

for num, model_name in enumerate(model_names):

name = model_name+"_ptm" if opt["use_ptm"] else model_name

key = f"{name}_seed_{seed}"

pbar.set_description(f'Running {key}')

# replace model parameters

params = data.get_model_haiku_params(name, params_loc)

for k in runner["model"].params.keys():

runner["model"].params[k] = params[k]

# predict

prediction_result, (r, t) = runner["model"].predict(processed_feature_dict, random_seed=seed)

outs[key] = parse_results(prediction_result, processed_feature_dict, r=r, t=t, num_res=num_res)

# cleanup

del prediction_result, params, r, t

# report

do_report(key)

pbar.update(n=1)

# cleanup

del processed_feature_dict

if subsample_msa: del feat

else:

# go through each model

for num, model_name in enumerate(model_names):

name = model_name+"_ptm" if opt["use_ptm"] else model_name

model_runner = prep_model_runner(opt, model_name=model_name, params_loc=params_loc)["model"]

# go through each random_seed

for seed in range(num_samples):

key = f"{name}_seed_{seed}"

pbar.set_description(f'Running {key}')

processed_feature_dict = model_runner.process_features(feature_dict, random_seed=seed)

# predict

prediction_result, (r, t) = model_runner.predict(processed_feature_dict, random_seed=seed)

outs[key] = parse_results(prediction_result, processed_feature_dict, r=r, t=t, num_res=num_res)

# cleanup

del processed_feature_dict, prediction_result, r, t

# report

do_report(key)

pbar.update(n=1)

# cleanup

del model_runner

# Find the best model according to the mean pLDDT.

model_rank = list(outs.keys())

model_rank = [model_rank[i] for i in np.argsort([outs[x][rank_by] for x in model_rank])[::-1]]

# Write out the prediction

for n,key in enumerate(model_rank):

prefix = f"rank_{n+1}_{key}"

pred_output_path = os.path.join(feature_dict["output_dir"],f'{prefix}_unrelaxed.pdb')

fig = cf.plot_protein(outs[key]["unrelaxed_protein"], Ls=feature_dict["Ls"], dpi=200)

plt.savefig(os.path.join(feature_dict["output_dir"],f'{prefix}.png'), bbox_inches = 'tight')

plt.close(fig)

pdb_lines = protein.to_pdb(outs[key]["unrelaxed_protein"])

with open(pred_output_path, 'w') as f:

f.write(pdb_lines)

tmp_pdb_path = os.path.join(feature_dict["output_dir"],f'unranked_{key}_unrelaxed.pdb')

if os.path.isfile(tmp_pdb_path):

os.remove(tmp_pdb_path)

############################################################

if verbose:

print(f"model rank based on {rank_by}")

for n,key in enumerate(model_rank):

print(f"rank_{n+1}_{key} {rank_by}:{outs[key][rank_by]:.2f}")