| import jax
|
| import jax.numpy as jnp
|
| import numpy as np
|
|
|
| from colabdesign.shared.utils import copy_dict
|
| from colabdesign.shared.model import soft_seq
|
| from colabdesign.af.alphafold.common import residue_constants
|
| from colabdesign.af.alphafold.model import model, config
|
|
|
|
|
|
|
|
|
| class _af_inputs:
|
|
|
| def _get_seq(self, inputs, aux, key=None):
|
| params, opt = inputs["params"], inputs["opt"]
|
| '''get sequence features'''
|
| seq = soft_seq(params["seq"], inputs["bias"], opt, key, num_seq=self._num,
|
| shuffle_first=self._args["shuffle_first"])
|
| seq = self._fix_pos(seq)
|
| aux.update({"seq":seq, "seq_pseudo":seq["pseudo"]})
|
|
|
|
|
| if self.protocol == "binder":
|
|
|
| seq_target = jax.nn.one_hot(inputs["batch"]["aatype"][:self._target_len],self._args["alphabet_size"])
|
| seq_target = jnp.broadcast_to(seq_target,(self._num, *seq_target.shape))
|
| seq = jax.tree_util.tree_map(lambda x:jnp.concatenate([seq_target,x],1), seq)
|
|
|
| if self.protocol in ["fixbb","hallucination","partial"] and self._args["copies"] > 1:
|
| seq = jax.tree_util.tree_map(lambda x:expand_copies(x, self._args["copies"], self._args["block_diag"]), seq)
|
|
|
| return seq
|
|
|
| def _fix_pos(self, seq, return_p=False):
|
| if "fix_pos" in self.opt:
|
| if "pos" in self.opt:
|
| seq_ref = jax.nn.one_hot(self._wt_aatype_sub,self._args["alphabet_size"])
|
| p = self.opt["pos"][self.opt["fix_pos"]]
|
| fix_seq = lambda x: x.at[...,p,:].set(seq_ref)
|
| else:
|
| seq_ref = jax.nn.one_hot(self._wt_aatype,self._args["alphabet_size"])
|
| p = self.opt["fix_pos"]
|
| fix_seq = lambda x: x.at[...,p,:].set(seq_ref[...,p,:])
|
| seq = jax.tree_util.tree_map(fix_seq, seq)
|
| if return_p: return seq, p
|
| return seq
|
|
|
| def _update_template(self, inputs, key):
|
| ''''dynamically update template features'''
|
| if "batch" in inputs:
|
| batch, opt = inputs["batch"], inputs["opt"]
|
|
|
|
|
| inputs["template_mask"] = inputs["template_mask"].at[0].set(1)
|
| L = batch["aatype"].shape[0]
|
|
|
|
|
| rm = jnp.broadcast_to(inputs.get("rm_template",False),L)
|
| rm_seq = jnp.where(rm,True,jnp.broadcast_to(inputs.get("rm_template_seq",True),L))
|
| rm_sc = jnp.where(rm_seq,True,jnp.broadcast_to(inputs.get("rm_template_sc",True),L))
|
|
|
|
|
| template_feats = {"template_aatype":jnp.where(rm_seq,21,batch["aatype"])}
|
|
|
| if "dgram" in batch:
|
|
|
| template_feats.update({"template_dgram":batch["dgram"]})
|
| nT,nL = inputs["template_aatype"].shape
|
| inputs["template_dgram"] = jnp.zeros((nT,nL,nL,39))
|
|
|
| if "all_atom_positions" in batch:
|
|
|
|
|
| cb, cb_mask = model.modules.pseudo_beta_fn(
|
| jnp.where(rm_seq,0,batch["aatype"]),
|
| batch["all_atom_positions"],
|
| batch["all_atom_mask"])
|
| template_feats.update({"template_pseudo_beta": cb,
|
| "template_pseudo_beta_mask": cb_mask,
|
| "template_all_atom_positions": batch["all_atom_positions"],
|
| "template_all_atom_mask": batch["all_atom_mask"]})
|
|
|
|
|
| if self.protocol == "partial":
|
| pos = opt["pos"]
|
| if self._args["repeat"] or self._args["homooligomer"]:
|
| C,L = self._args["copies"], self._len
|
| pos = (jnp.repeat(pos,C).reshape(-1,C) + jnp.arange(C) * L).T.flatten()
|
|
|
| for k,v in template_feats.items():
|
| if self.protocol == "partial":
|
| if k in ["template_dgram"]:
|
| inputs[k] = inputs[k].at[0,pos[:,None],pos[None,:]].set(v)
|
| else:
|
| inputs[k] = inputs[k].at[0,pos].set(v)
|
| else:
|
| inputs[k] = inputs[k].at[0].set(v)
|
|
|
|
|
| if k in ["template_all_atom_mask"]:
|
| if self.protocol == "partial":
|
| inputs[k] = inputs[k].at[:,pos,5:].set(jnp.where(rm_sc[:,None],0,inputs[k][:,pos,5:]))
|
| inputs[k] = inputs[k].at[:,pos].set(jnp.where(rm[:,None],0,inputs[k][:,pos]))
|
| else:
|
| inputs[k] = inputs[k].at[...,5:].set(jnp.where(rm_sc[:,None],0,inputs[k][...,5:]))
|
| inputs[k] = jnp.where(rm[:,None],0,inputs[k])
|
|
|
| def update_seq(seq, inputs, seq_1hot=None, seq_pssm=None, mlm=None):
|
| '''update the sequence features'''
|
|
|
| if seq_1hot is None: seq_1hot = seq
|
| if seq_pssm is None: seq_pssm = seq
|
| target_feat = seq_1hot[0,:,:20]
|
|
|
| seq_1hot = jnp.pad(seq_1hot,[[0,0],[0,0],[0,22-seq_1hot.shape[-1]]])
|
| seq_pssm = jnp.pad(seq_pssm,[[0,0],[0,0],[0,22-seq_pssm.shape[-1]]])
|
| msa_feat = jnp.zeros_like(inputs["msa_feat"]).at[...,0:22].set(seq_1hot).at[...,25:47].set(seq_pssm)
|
|
|
|
|
| if mlm is not None:
|
| X = jax.nn.one_hot(22,23)
|
| X = jnp.zeros(msa_feat.shape[-1]).at[...,:23].set(X).at[...,25:48].set(X)
|
| msa_feat = jnp.where(mlm[...,None],X,msa_feat)
|
|
|
| inputs.update({"msa_feat":msa_feat, "target_feat":target_feat})
|
|
|
| def update_aatype(aatype, inputs):
|
| r = residue_constants
|
| a = {"atom14_atom_exists":r.restype_atom14_mask,
|
| "atom37_atom_exists":r.restype_atom37_mask,
|
| "residx_atom14_to_atom37":r.restype_atom14_to_atom37,
|
| "residx_atom37_to_atom14":r.restype_atom37_to_atom14}
|
| mask = inputs["seq_mask"][:,None]
|
| inputs.update(jax.tree_util.tree_map(lambda x:jnp.where(mask,jnp.asarray(x)[aatype],0),a))
|
| inputs["aatype"] = aatype
|
|
|
| def expand_copies(x, copies, block_diag=True):
|
| '''
|
| given msa (N,L,20) expand to (1+N*copies,L*copies,22) if block_diag else (N,L*copies,22)
|
| '''
|
| if x.shape[-1] < 22:
|
| x = jnp.pad(x,[[0,0],[0,0],[0,22-x.shape[-1]]])
|
| x = jnp.tile(x,[1,copies,1])
|
| if copies > 1 and block_diag:
|
| L = x.shape[1]
|
| sub_L = L // copies
|
| y = x.reshape((-1,1,copies,sub_L,22))
|
| block_diag_mask = jnp.expand_dims(jnp.eye(copies),(0,3,4))
|
| seq = block_diag_mask * y
|
| gap_seq = (1-block_diag_mask) * jax.nn.one_hot(jnp.repeat(21,sub_L),22)
|
| y = (seq + gap_seq).swapaxes(0,1).reshape(-1,L,22)
|
| return jnp.concatenate([x[:1],y],0)
|
| else:
|
| return x |
