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()) ################################################# # AF_PREP - input prep functions ################################################# 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) --------------------------------------------------- ''' # prep features 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] # feat dims num_seq = self._num res_idx = self._pdb["residue_index"] # get [pos]itions of interests 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(",")) # repeat/homo-oligomeric support 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"] # configure input features 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)) # configure options/weights 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] # configure template [opt]ions 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 --------------------------------------------------- ''' # define num copies (for repeats/ homo-oligomers) 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}) # prep features self._len = length # set weights 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) # configure input features 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, # if binder_chain is defined 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}) # get pdb info target_chain = kwargs.pop("chain",target_chain) # backward comp 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] # gather hotspot info if hotspot is not None: self.opt["hotspot"] = prep_pos(hotspot, **self._pdb["idx"])["pos"] if redesign: # binder 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: # binder hallucination 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}) # configure input features 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)) # configure template rm masks (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 # set template [opt]ions 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) --------------------------------------------------- ''' # prep features 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] # feat dims num_seq = self._num res_idx = np.arange(self._len) # get [pos]itions of interests 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(",")) # repeat/homo-oligomeric support 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"]] # repeat positions across copies 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 # configure input features 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)) # configure options/weights 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"]] # configure sidechains 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) ####################### # utils ####################### 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 = [] # go through each defined chain 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) # add in missing cb (in the case of glycine) 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: # pad values 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)) # concatenate chains o = jax.tree_util.tree_map(lambda *x:np.concatenate(x,0),*o) # save original residue and chain index 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''' # decide what atoms to exclude for each residue type 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 # collect atom indices 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)), # 23 = one_hot -> (20, UNK, GAP, MASK) # 1 = has deletion # 1 = deletion_value # 23 = profile # 1 = deletion_mean_value '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), # for template inputs '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)), # for alphafold-multimer '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}