| import jax
|
| import jax.numpy as jnp
|
| import numpy as np
|
| import re
|
|
|
| from colabdesign.af.alphafold.data import pipeline, prep_inputs
|
| from colabdesign.af.alphafold.common import protein, residue_constants
|
| from colabdesign.af.alphafold.model.tf import shape_placeholders
|
| from colabdesign.af.alphafold.model import config
|
|
|
|
|
| from colabdesign.shared.protein import _np_get_cb, pdb_to_string
|
| from colabdesign.shared.prep import prep_pos
|
| from colabdesign.shared.utils import copy_dict
|
| from colabdesign.shared.model import order_aa
|
|
|
| resname_to_idx = residue_constants.resname_to_idx
|
| idx_to_resname = dict((v,k) for k,v in resname_to_idx.items())
|
|
|
|
|
|
|
|
|
| class _af_prep:
|
|
|
| def _prep_model(self, **kwargs):
|
| '''prep model'''
|
| if not hasattr(self,"_model") or self._cfg != self._model["runner"].config:
|
| self._cfg.model.global_config.subbatch_size = None
|
| self._model = self._get_model(self._cfg)
|
| if sum(self._lengths) > 384:
|
| self._cfg.model.global_config.subbatch_size = 4
|
| self._model["fn"] = self._get_model(self._cfg)["fn"]
|
|
|
| self._opt = copy_dict(self.opt)
|
| self.restart(**kwargs)
|
|
|
| def _prep_features(self, num_res, num_seq=None, num_templates=1):
|
| '''process features'''
|
| if num_seq is None: num_seq = self._num
|
| return prep_input_features(L=num_res, N=num_seq, T=num_templates)
|
|
|
| def _prep_fixbb(self, pdb_filename, chain=None,
|
| copies=1, repeat=False, homooligomer=False,
|
| rm_template=False,
|
| rm_template_seq=True,
|
| rm_template_sc=True,
|
| rm_template_ic=False,
|
| fix_pos=None, ignore_missing=True, **kwargs):
|
| '''
|
| prep inputs for fixed backbone design
|
| ---------------------------------------------------
|
| if copies > 1:
|
| -homooligomer=True - input pdb chains are parsed as homo-oligomeric units
|
| -repeat=True - tie the repeating sequence within single chain
|
| -rm_template_seq - if template is defined, remove information about template sequence
|
| -fix_pos="1,2-10" - specify which positions to keep fixed in the sequence
|
| note: supervised loss is applied to all positions, use "partial"
|
| protocol to apply supervised loss to only subset of positions
|
| -ignore_missing=True - skip positions that have missing density (no CA coordinate)
|
| ---------------------------------------------------
|
| '''
|
|
|
| self._pdb = prep_pdb(pdb_filename, chain=chain, ignore_missing=ignore_missing,
|
| offsets=kwargs.pop("pdb_offsets",None),
|
| lengths=kwargs.pop("pdb_lengths",None))
|
|
|
| self._len = self._pdb["residue_index"].shape[0]
|
| self._lengths = [self._len]
|
|
|
|
|
| num_seq = self._num
|
| res_idx = self._pdb["residue_index"]
|
|
|
|
|
| if fix_pos is not None and fix_pos != "":
|
| self._pos_info = prep_pos(fix_pos, **self._pdb["idx"])
|
| self.opt["fix_pos"] = self._pos_info["pos"]
|
|
|
| if homooligomer and chain is not None and copies == 1:
|
| copies = len(chain.split(","))
|
|
|
|
|
| if copies > 1:
|
|
|
| if repeat or homooligomer:
|
| self._len = self._len // copies
|
| if "fix_pos" in self.opt:
|
| self.opt["fix_pos"] = self.opt["fix_pos"][self.opt["fix_pos"] < self._len]
|
|
|
| if repeat:
|
| self._lengths = [self._len * copies]
|
| block_diag = False
|
|
|
| else:
|
| self._lengths = [self._len] * copies
|
| block_diag = not self._args["use_multimer"]
|
|
|
| res_idx = repeat_idx(res_idx[:self._len], copies)
|
| num_seq = (self._num * copies + 1) if block_diag else self._num
|
| self.opt["weights"].update({"i_pae":0.0, "i_con":0.0})
|
|
|
| self._args.update({"copies":copies, "repeat":repeat, "homooligomer":homooligomer, "block_diag":block_diag})
|
| homooligomer = not repeat
|
| else:
|
| self._lengths = self._pdb["lengths"]
|
|
|
|
|
| self._inputs = self._prep_features(num_res=sum(self._lengths), num_seq=num_seq)
|
| self._inputs["residue_index"] = res_idx
|
| self._inputs["batch"] = make_fixed_size(self._pdb["batch"], num_res=sum(self._lengths))
|
| self._inputs.update(get_multi_id(self._lengths, homooligomer=homooligomer))
|
|
|
|
|
| self.opt["weights"].update({"dgram_cce":1.0, "rmsd":0.0, "fape":0.0, "con":0.0})
|
| self._wt_aatype = self._inputs["batch"]["aatype"][:self._len]
|
|
|
|
|
| rm,L = {},sum(self._lengths)
|
| for n,x in {"rm_template": rm_template,
|
| "rm_template_seq":rm_template_seq,
|
| "rm_template_sc": rm_template_sc}.items():
|
| rm[n] = np.full(L,False)
|
| if isinstance(x,str):
|
| rm[n][prep_pos(x,**self._pdb["idx"])["pos"]] = True
|
| else:
|
| rm[n][:] = x
|
| self.opt["template"]["rm_ic"] = rm_template_ic
|
| self._inputs.update(rm)
|
|
|
| self._prep_model(**kwargs)
|
|
|
| def _prep_hallucination(self, length=100, copies=1, repeat=False, **kwargs):
|
| '''
|
| prep inputs for hallucination
|
| ---------------------------------------------------
|
| if copies > 1:
|
| -repeat=True - tie the repeating sequence within single chain
|
| ---------------------------------------------------
|
| '''
|
|
|
|
|
| if not repeat and copies > 1 and not self._args["use_multimer"]:
|
| (num_seq, block_diag) = (self._num * copies + 1, True)
|
| else:
|
| (num_seq, block_diag) = (self._num, False)
|
|
|
| self._args.update({"repeat":repeat,"block_diag":block_diag,"copies":copies})
|
|
|
|
|
| self._len = length
|
|
|
|
|
| self.opt["weights"].update({"con":1.0})
|
| if copies > 1:
|
| if repeat:
|
| offset = 1
|
| self._lengths = [self._len * copies]
|
| self._args["repeat"] = True
|
| else:
|
| offset = 50
|
| self._lengths = [self._len] * copies
|
| self.opt["weights"].update({"i_pae":0.0, "i_con":1.0})
|
| self._args["homooligomer"] = True
|
| res_idx = repeat_idx(np.arange(length), copies, offset=offset)
|
| else:
|
| self._lengths = [self._len]
|
| res_idx = np.arange(length)
|
|
|
|
|
| self._inputs = self._prep_features(num_res=sum(self._lengths), num_seq=num_seq)
|
| self._inputs["residue_index"] = res_idx
|
| self._inputs.update(get_multi_id(self._lengths, homooligomer=True))
|
|
|
| self._prep_model(**kwargs)
|
|
|
| def _prep_binder(self, pdb_filename,
|
| target_chain="A", binder_len=50,
|
| rm_target = False,
|
| rm_target_seq = False,
|
| rm_target_sc = False,
|
|
|
|
|
| binder_chain=None,
|
| rm_binder=True,
|
| rm_binder_seq=True,
|
| rm_binder_sc=True,
|
| rm_template_ic=False,
|
|
|
| hotspot=None, ignore_missing=True, **kwargs):
|
| '''
|
| prep inputs for binder design
|
| ---------------------------------------------------
|
| -binder_len = length of binder to hallucinate (option ignored if binder_chain is defined)
|
| -binder_chain = chain of binder to redesign
|
| -use_binder_template = use binder coordinates as template input
|
| -rm_template_ic = use target and binder coordinates as seperate template inputs
|
| -hotspot = define position/hotspots on target
|
| -rm_[binder/target]_seq = remove sequence info from template
|
| -rm_[binder/target]_sc = remove sidechain info from template
|
| -ignore_missing=True - skip positions that have missing density (no CA coordinate)
|
| ---------------------------------------------------
|
| '''
|
| redesign = binder_chain is not None
|
| rm_binder = not kwargs.pop("use_binder_template", not rm_binder)
|
|
|
| self._args.update({"redesign":redesign})
|
|
|
|
|
| target_chain = kwargs.pop("chain",target_chain)
|
| chains = f"{target_chain},{binder_chain}" if redesign else target_chain
|
| im = [True] * len(target_chain.split(","))
|
| if redesign: im += [ignore_missing] * len(binder_chain.split(","))
|
|
|
| self._pdb = prep_pdb(pdb_filename, chain=chains, ignore_missing=im)
|
| res_idx = self._pdb["residue_index"]
|
|
|
| if redesign:
|
| self._target_len = sum([(self._pdb["idx"]["chain"] == c).sum() for c in target_chain.split(",")])
|
| self._binder_len = sum([(self._pdb["idx"]["chain"] == c).sum() for c in binder_chain.split(",")])
|
| else:
|
| self._target_len = self._pdb["residue_index"].shape[0]
|
| self._binder_len = binder_len
|
| res_idx = np.append(res_idx, res_idx[-1] + np.arange(binder_len) + 50)
|
|
|
| self._len = self._binder_len
|
| self._lengths = [self._target_len, self._binder_len]
|
|
|
|
|
| if hotspot is not None:
|
| self.opt["hotspot"] = prep_pos(hotspot, **self._pdb["idx"])["pos"]
|
|
|
| if redesign:
|
|
|
| self._wt_aatype = self._pdb["batch"]["aatype"][self._target_len:]
|
| self.opt["weights"].update({"dgram_cce":1.0, "rmsd":0.0, "fape":0.0,
|
| "con":0.0, "i_con":0.0, "i_pae":0.0})
|
| else:
|
|
|
| self._pdb["batch"] = make_fixed_size(self._pdb["batch"], num_res=sum(self._lengths))
|
| self.opt["weights"].update({"plddt":0.1, "con":0.0, "i_con":1.0, "i_pae":0.0})
|
|
|
|
|
| self._inputs = self._prep_features(num_res=sum(self._lengths), num_seq=1)
|
| self._inputs["residue_index"] = res_idx
|
| self._inputs["batch"] = self._pdb["batch"]
|
| self._inputs.update(get_multi_id(self._lengths))
|
|
|
|
|
| (T,L,rm) = (self._lengths[0],sum(self._lengths),{})
|
| rm_opt = {
|
| "rm_template": {"target":rm_target, "binder":rm_binder},
|
| "rm_template_seq":{"target":rm_target_seq,"binder":rm_binder_seq},
|
| "rm_template_sc": {"target":rm_target_sc, "binder":rm_binder_sc}
|
| }
|
| for n,x in rm_opt.items():
|
| rm[n] = np.full(L,False)
|
| for m,y in x.items():
|
| if isinstance(y,str):
|
| rm[n][prep_pos(y,**self._pdb["idx"])["pos"]] = True
|
| else:
|
| if m == "target": rm[n][:T] = y
|
| if m == "binder": rm[n][T:] = y
|
|
|
|
|
| self.opt["template"]["rm_ic"] = rm_template_ic
|
| self._inputs.update(rm)
|
|
|
| self._prep_model(**kwargs)
|
|
|
| def _prep_partial(self, pdb_filename, chain=None, length=None,
|
| copies=1, repeat=False, homooligomer=False,
|
| pos=None, fix_pos=None, use_sidechains=False, atoms_to_exclude=None,
|
| rm_template=False,
|
| rm_template_seq=False,
|
| rm_template_sc=False,
|
| rm_template_ic=False,
|
| ignore_missing=True, **kwargs):
|
| '''
|
| prep input for partial hallucination
|
| ---------------------------------------------------
|
| -length=100 - total length of protein (if different from input PDB)
|
| -pos="1,2-10" - specify which positions to apply supervised loss to
|
| -use_sidechains=True - add a sidechain supervised loss to the specified positions
|
| -atoms_to_exclude=["N","C","O"] (for sc_rmsd loss, specify which atoms to exclude)
|
| -rm_template_seq - if template is defined, remove information about template sequence
|
| -ignore_missing=True - skip positions that have missing density (no CA coordinate)
|
| ---------------------------------------------------
|
| '''
|
|
|
| self._pdb = prep_pdb(pdb_filename, chain=chain, ignore_missing=ignore_missing,
|
| offsets=kwargs.pop("pdb_offsets",None),
|
| lengths=kwargs.pop("pdb_lengths",None))
|
|
|
| self._pdb["len"] = sum(self._pdb["lengths"])
|
|
|
| self._len = self._pdb["len"] if length is None else length
|
| self._lengths = [self._len]
|
|
|
|
|
| num_seq = self._num
|
| res_idx = np.arange(self._len)
|
|
|
|
|
| if pos is None:
|
| self.opt["pos"] = self._pdb["pos"] = np.arange(self._pdb["len"])
|
| self._pos_info = {"length":np.array([self._pdb["len"]]), "pos":self._pdb["pos"]}
|
| else:
|
| self._pos_info = prep_pos(pos, **self._pdb["idx"])
|
| self.opt["pos"] = self._pdb["pos"] = self._pos_info["pos"]
|
|
|
| if homooligomer and chain is not None and copies == 1:
|
| copies = len(chain.split(","))
|
|
|
|
|
| if copies > 1:
|
|
|
| if repeat or homooligomer:
|
| self._len = self._len // copies
|
| self._pdb["len"] = self._pdb["len"] // copies
|
| self.opt["pos"] = self._pdb["pos"][self._pdb["pos"] < self._pdb["len"]]
|
|
|
|
|
| self._pdb["pos"] = repeat_pos(self.opt["pos"], copies, self._pdb["len"])
|
|
|
| if repeat:
|
| self._lengths = [self._len * copies]
|
| block_diag = False
|
|
|
| else:
|
| self._lengths = [self._len] * copies
|
| block_diag = not self._args["use_multimer"]
|
|
|
| num_seq = (self._num * copies + 1) if block_diag else self._num
|
| res_idx = repeat_idx(np.arange(self._len), copies)
|
|
|
| self.opt["weights"].update({"i_pae":0.0, "i_con":1.0})
|
|
|
| self._args.update({"copies":copies, "repeat":repeat, "homooligomer":homooligomer, "block_diag":block_diag})
|
| homooligomer = not repeat
|
|
|
|
|
| self._inputs = self._prep_features(num_res=sum(self._lengths), num_seq=num_seq)
|
| self._inputs["residue_index"] = res_idx
|
| self._inputs["batch"] = jax.tree_util.tree_map(lambda x:x[self._pdb["pos"]], self._pdb["batch"])
|
| self._inputs.update(get_multi_id(self._lengths, homooligomer=homooligomer))
|
|
|
|
|
| self.opt["weights"].update({"dgram_cce":1.0, "rmsd":0.0, "fape":0.0, "con":1.0})
|
| self._wt_aatype = self._pdb["batch"]["aatype"][self.opt["pos"]]
|
|
|
|
|
| self._args["use_sidechains"] = use_sidechains
|
| if use_sidechains:
|
| self._sc = {"batch":prep_inputs.make_atom14_positions(self._inputs["batch"]),
|
| "pos":get_sc_pos(self._wt_aatype, atoms_to_exclude)}
|
| self.opt["weights"].update({"sc_rmsd":0.1, "sc_fape":0.1})
|
| self.opt["fix_pos"] = np.arange(self.opt["pos"].shape[0])
|
| self._wt_aatype_sub = self._wt_aatype
|
|
|
| elif fix_pos is not None and fix_pos != "":
|
| sub_fix_pos = []
|
| sub_i = []
|
| pos = self.opt["pos"].tolist()
|
| for i in prep_pos(fix_pos, **self._pdb["idx"])["pos"]:
|
| if i in pos:
|
| sub_i.append(i)
|
| sub_fix_pos.append(pos.index(i))
|
| self.opt["fix_pos"] = np.array(sub_fix_pos)
|
| self._wt_aatype_sub = self._pdb["batch"]["aatype"][sub_i]
|
|
|
| elif kwargs.pop("fix_seq",False):
|
| self.opt["fix_pos"] = np.arange(self.opt["pos"].shape[0])
|
| self._wt_aatype_sub = self._wt_aatype
|
|
|
| self.opt["template"].update({"rm_ic":rm_template_ic})
|
| self._inputs.update({"rm_template": rm_template,
|
| "rm_template_seq": rm_template_seq,
|
| "rm_template_sc": rm_template_sc})
|
|
|
| self._prep_model(**kwargs)
|
|
|
|
|
|
|
|
|
| def repeat_idx(idx, copies=1, offset=50):
|
| idx_offset = np.repeat(np.cumsum([0]+[idx[-1]+offset]*(copies-1)),len(idx))
|
| return np.tile(idx,copies) + idx_offset
|
|
|
| def repeat_pos(pos, copies, length):
|
| return (np.repeat(pos,copies).reshape(-1,copies) + np.arange(copies) * length).T.flatten()
|
|
|
| def prep_pdb(pdb_filename, chain=None,
|
| offsets=None, lengths=None,
|
| ignore_missing=False):
|
| '''extract features from pdb'''
|
|
|
| def add_cb(batch):
|
| '''add missing CB atoms based on N,CA,C'''
|
| p,m = batch["all_atom_positions"], batch["all_atom_mask"]
|
| atom_idx = residue_constants.atom_order
|
| atoms = {k:p[...,atom_idx[k],:] for k in ["N","CA","C"]}
|
| cb = atom_idx["CB"]
|
| cb_atoms = _np_get_cb(**atoms, use_jax=False)
|
| cb_mask = np.prod([m[...,atom_idx[k]] for k in ["N","CA","C"]],0)
|
| batch["all_atom_positions"][...,cb,:] = np.where(m[:,cb,None], p[:,cb,:], cb_atoms)
|
| batch["all_atom_mask"][...,cb] = (m[:,cb] + cb_mask) > 0
|
| return {"atoms":batch["all_atom_positions"][:,cb],"mask":cb_mask}
|
|
|
| if isinstance(chain,str) and "," in chain:
|
| chains = chain.split(",")
|
| elif not isinstance(chain,list):
|
| chains = [chain]
|
|
|
| o,last = [],0
|
| residue_idx, chain_idx = [],[]
|
| full_lengths = []
|
|
|
|
|
| for n,chain in enumerate(chains):
|
| pdb_str = pdb_to_string(pdb_filename, chains=chain, models=[1])
|
| protein_obj = protein.from_pdb_string(pdb_str, chain_id=chain)
|
| batch = {'aatype': protein_obj.aatype,
|
| 'all_atom_positions': protein_obj.atom_positions,
|
| 'all_atom_mask': protein_obj.atom_mask,
|
| 'residue_index': protein_obj.residue_index}
|
|
|
| cb_feat = add_cb(batch)
|
|
|
| im = ignore_missing[n] if isinstance(ignore_missing,list) else ignore_missing
|
| if im:
|
| r = batch["all_atom_mask"][:,0] == 1
|
| batch = jax.tree_util.tree_map(lambda x:x[r], batch)
|
| residue_index = batch["residue_index"] + last
|
|
|
| else:
|
|
|
| offset = 0 if offsets is None else (offsets[n] if isinstance(offsets,list) else offsets)
|
| r = offset + (protein_obj.residue_index - protein_obj.residue_index.min())
|
| length = (r.max()+1) if lengths is None else (lengths[n] if isinstance(lengths,list) else lengths)
|
| def scatter(x, value=0):
|
| shape = (length,) + x.shape[1:]
|
| y = np.full(shape, value, dtype=x.dtype)
|
| y[r] = x
|
| return y
|
|
|
| batch = {"aatype":scatter(batch["aatype"],-1),
|
| "all_atom_positions":scatter(batch["all_atom_positions"]),
|
| "all_atom_mask":scatter(batch["all_atom_mask"]),
|
| "residue_index":scatter(batch["residue_index"],-1)}
|
|
|
| residue_index = np.arange(length) + last
|
|
|
| last = residue_index[-1] + 50
|
| o.append({"batch":batch,
|
| "residue_index": residue_index,
|
| "cb_feat":cb_feat})
|
|
|
| residue_idx.append(batch.pop("residue_index"))
|
| chain_idx.append([chain] * len(residue_idx[-1]))
|
| full_lengths.append(len(residue_index))
|
|
|
|
|
| o = jax.tree_util.tree_map(lambda *x:np.concatenate(x,0),*o)
|
|
|
|
|
| o["idx"] = {"residue":np.concatenate(residue_idx), "chain":np.concatenate(chain_idx)}
|
| o["lengths"] = full_lengths
|
| return o
|
|
|
| def make_fixed_size(feat, num_res, num_seq=1, num_templates=1):
|
| '''pad input features'''
|
| shape_schema = {k:v for k,v in config.CONFIG.data.eval.feat.items()}
|
|
|
| pad_size_map = {
|
| shape_placeholders.NUM_RES: num_res,
|
| shape_placeholders.NUM_MSA_SEQ: num_seq,
|
| shape_placeholders.NUM_EXTRA_SEQ: 1,
|
| shape_placeholders.NUM_TEMPLATES: num_templates
|
| }
|
| for k,v in feat.items():
|
| if k == "batch":
|
| feat[k] = make_fixed_size(v, num_res)
|
| else:
|
| 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 - v.shape[i]) for i, p in enumerate(pad_size)]
|
| feat[k] = np.pad(v, padding)
|
| return feat
|
|
|
| def get_sc_pos(aa_ident, atoms_to_exclude=None):
|
| '''get sidechain indices/weights for all_atom14_positions'''
|
|
|
|
|
| a2e = {}
|
| for r in resname_to_idx:
|
| if isinstance(atoms_to_exclude,dict):
|
| a2e[r] = atoms_to_exclude.get(r,atoms_to_exclude.get("ALL",["N","C","O"]))
|
| else:
|
| a2e[r] = ["N","C","O"] if atoms_to_exclude is None else atoms_to_exclude
|
|
|
|
|
| pos,pos_alt = [],[]
|
| N,N_non_amb = [],[]
|
| for n,a in enumerate(aa_ident):
|
| aa = idx_to_resname[a]
|
| atoms = set(residue_constants.residue_atoms[aa])
|
| atoms14 = residue_constants.restype_name_to_atom14_names[aa]
|
| swaps = residue_constants.residue_atom_renaming_swaps.get(aa,{})
|
| swaps.update({v:k for k,v in swaps.items()})
|
| for atom in atoms.difference(a2e[aa]):
|
| pos.append(n * 14 + atoms14.index(atom))
|
| if atom in swaps:
|
| pos_alt.append(n * 14 + atoms14.index(swaps[atom]))
|
| else:
|
| pos_alt.append(pos[-1])
|
| N_non_amb.append(n)
|
| N.append(n)
|
|
|
| pos, pos_alt = np.asarray(pos), np.asarray(pos_alt)
|
| non_amb = pos == pos_alt
|
| N, N_non_amb = np.asarray(N), np.asarray(N_non_amb)
|
| w = np.array([1/(n == N).sum() for n in N])
|
| w_na = np.array([1/(n == N_non_amb).sum() for n in N_non_amb])
|
| w, w_na = w/w.sum(), w_na/w_na.sum()
|
| return {"pos":pos, "pos_alt":pos_alt, "non_amb":non_amb,
|
| "weight":w, "weight_non_amb":w_na[:,None]}
|
|
|
| def prep_input_features(L, N=1, T=1, eN=1):
|
| '''
|
| given [L]ength, [N]umber of sequences and number of [T]emplates
|
| return dictionary of blank features
|
| '''
|
| inputs = {'aatype': np.zeros(L,int),
|
| 'target_feat': np.zeros((L,20)),
|
| 'msa_feat': np.zeros((N,L,49)),
|
|
|
|
|
|
|
|
|
|
|
|
|
| 'seq_mask': np.ones(L),
|
| 'msa_mask': np.ones((N,L)),
|
| 'msa_row_mask': np.ones(N),
|
| 'atom14_atom_exists': np.zeros((L,14)),
|
| 'atom37_atom_exists': np.zeros((L,37)),
|
| 'residx_atom14_to_atom37': np.zeros((L,14),int),
|
| 'residx_atom37_to_atom14': np.zeros((L,37),int),
|
| 'residue_index': np.arange(L),
|
| 'extra_deletion_value': np.zeros((eN,L)),
|
| 'extra_has_deletion': np.zeros((eN,L)),
|
| 'extra_msa': np.zeros((eN,L),int),
|
| 'extra_msa_mask': np.zeros((eN,L)),
|
| 'extra_msa_row_mask': np.zeros(eN),
|
|
|
|
|
| 'template_aatype': np.zeros((T,L),int),
|
| 'template_all_atom_mask': np.zeros((T,L,37)),
|
| 'template_all_atom_positions': np.zeros((T,L,37,3)),
|
| 'template_mask': np.zeros(T),
|
| 'template_pseudo_beta': np.zeros((T,L,3)),
|
| 'template_pseudo_beta_mask': np.zeros((T,L)),
|
|
|
|
|
| 'asym_id': np.zeros(L),
|
| 'sym_id': np.zeros(L),
|
| 'entity_id': np.zeros(L),
|
| 'all_atom_positions': np.zeros((N,37,3))}
|
| return inputs
|
|
|
| def get_multi_id(lengths, homooligomer=False):
|
| '''set info for alphafold-multimer'''
|
| i = np.concatenate([[n]*l for n,l in enumerate(lengths)])
|
| if homooligomer:
|
| return {"asym_id":i, "sym_id":i, "entity_id":np.zeros_like(i)}
|
| else:
|
| return {"asym_id":i, "sym_id":i, "entity_id":i} |
