model/PXDesignBench/ColabDesign/colabdesign/af/inputs.py

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

############################################################################
# AF_INPUTS - functions for modifying inputs before passing to alphafold
############################################################################
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"]})
    
    # protocol specific modifications to seq features
    if self.protocol == "binder":
      # concatenate target and binder sequence
      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"]

      # enable templates
      inputs["template_mask"] = inputs["template_mask"].at[0].set(1)
      L = batch["aatype"].shape[0]
      
      # decide which position to remove sequence and/or sidechains
      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))
                          
      # define template features
      template_feats = {"template_aatype":jnp.where(rm_seq,21,batch["aatype"])}

      if "dgram" in batch:
        # use dgram from batch if provided
        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:
        # get pseudo-carbon-beta coordinates (carbon-alpha for glycine)
        # aatype = is used to define template's CB coordinates (CA in case of glycine)
        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"]})

      # inject template features
      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)
        
        # remove sidechains (mask anything beyond CB)
        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)

  # masked language modeling (randomly mask positions)
  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