import functools import haiku as hk import jax import jax.numpy as jnp import numpy as np import itertools import joblib from colabdesign.shared.prng import SafeKey from .utils import gather_edges, gather_nodes, cat_neighbors_nodes, scatter, get_ar_mask from .sample import mpnn_sample Gelu = functools.partial(jax.nn.gelu, approximate=False) class dropout_cust(hk.Module): def __init__(self, rate) -> None: super().__init__() self.rate = rate self.safe_key = SafeKey(hk.next_rng_key()) def __call__(self, x): self.safe_key, use_key = self.safe_key.split() return hk.dropout(use_key.get(), self.rate, x) class EncLayer(hk.Module): def __init__(self, num_hidden, num_in, dropout=0.1, num_heads=None, scale=30, name=None): super(EncLayer, self).__init__() self.num_hidden = num_hidden self.num_in = num_in self.scale = scale self.safe_key = SafeKey(hk.next_rng_key()) self.dropout1 = dropout_cust(dropout) self.dropout2 = dropout_cust(dropout) self.dropout3 = dropout_cust(dropout) self.norm1 = hk.LayerNorm(-1, create_scale=True, create_offset=True, name=name + '_norm1') self.norm2 = hk.LayerNorm(-1, create_scale=True, create_offset=True, name=name + '_norm2') self.norm3 = hk.LayerNorm(-1, create_scale=True, create_offset=True, name=name + '_norm3') self.W1 = hk.Linear(num_hidden, with_bias=True, name=name + '_W1') self.W2 = hk.Linear(num_hidden, with_bias=True, name=name + '_W2') self.W3 = hk.Linear(num_hidden, with_bias=True, name=name + '_W3') self.W11 = hk.Linear(num_hidden, with_bias=True, name=name + '_W11') self.W12 = hk.Linear(num_hidden, with_bias=True, name=name + '_W12') self.W13 = hk.Linear(num_hidden, with_bias=True, name=name + '_W13') self.act = Gelu self.dense = PositionWiseFeedForward(num_hidden, num_hidden * 4, name=name + '_dense') def __call__(self, h_V, h_E, E_idx, mask_V=None, mask_attend=None): """ Parallel computation of full transformer layer """ h_EV = cat_neighbors_nodes(h_V, h_E, E_idx) h_V_expand = jnp.tile(jnp.expand_dims(h_V, -2),[1, 1, h_EV.shape[-2], 1]) h_EV = jnp.concatenate([h_V_expand, h_EV], -1) h_message = self.W3(self.act(self.W2(self.act(self.W1(h_EV))))) if mask_attend is not None: h_message = jnp.expand_dims(mask_attend, -1)* h_message dh = jnp.sum(h_message, -2) / self.scale h_V = self.norm1(h_V + self.dropout1(dh)) dh = self.dense(h_V) h_V = self.norm2(h_V + self.dropout2(dh)) if mask_V is not None: mask_V = jnp.expand_dims(mask_V, -1) h_V = mask_V * h_V h_EV = cat_neighbors_nodes(h_V, h_E, E_idx) h_V_expand = jnp.tile(jnp.expand_dims(h_V, -2),[1, 1, h_EV.shape[-2], 1]) h_EV = jnp.concatenate([h_V_expand, h_EV], -1) h_message = self.W13(self.act(self.W12(self.act(self.W11(h_EV))))) h_E = self.norm3(h_E + self.dropout3(h_message)) return h_V, h_E class DecLayer(hk.Module): def __init__(self, num_hidden, num_in, dropout=0.1, num_heads=None, scale=30, name=None): super(DecLayer, self).__init__() self.num_hidden = num_hidden self.num_in = num_in self.scale = scale self.dropout1 = dropout_cust(dropout) self.dropout2 = dropout_cust(dropout) self.norm1 = hk.LayerNorm(-1, create_scale=True, create_offset=True, name=name + '_norm1') self.norm2 = hk.LayerNorm(-1, create_scale=True, create_offset=True, name=name + '_norm2') self.W1 = hk.Linear(num_hidden, with_bias=True, name=name + '_W1') self.W2 = hk.Linear(num_hidden, with_bias=True, name=name + '_W2') self.W3 = hk.Linear(num_hidden, with_bias=True, name=name + '_W3') self.act = Gelu self.dense = PositionWiseFeedForward(num_hidden, num_hidden * 4, name=name + '_dense') def __call__(self, h_V, h_E, mask_V=None, mask_attend=None): """ Parallel computation of full transformer layer """ # Concatenate h_V_i to h_E_ij h_V_expand = jnp.tile(jnp.expand_dims(h_V, -2),[1, 1, h_E.shape[-2], 1]) h_EV = jnp.concatenate([h_V_expand, h_E], -1) h_message = self.W3(self.act(self.W2(self.act(self.W1(h_EV))))) if mask_attend is not None: h_message = jnp.expand_dims(mask_attend, -1) * h_message dh = jnp.sum(h_message, -2) / self.scale h_V = self.norm1(h_V + self.dropout1(dh)) # Position-wise feedforward dh = self.dense(h_V) h_V = self.norm2(h_V + self.dropout2(dh)) if mask_V is not None: mask_V = jnp.expand_dims(mask_V, -1) h_V = mask_V * h_V return h_V class PositionWiseFeedForward(hk.Module): def __init__(self, num_hidden, num_ff, name=None): super(PositionWiseFeedForward, self).__init__() self.W_in = hk.Linear(num_ff, with_bias=True, name=name + '_W_in') self.W_out = hk.Linear(num_hidden, with_bias=True, name=name + '_W_out') self.act = Gelu def __call__(self, h_V): h = self.act(self.W_in(h_V), approximate=False) h = self.W_out(h) return h class PositionalEncodings(hk.Module): def __init__(self, num_embeddings, max_relative_feature=32): super(PositionalEncodings, self).__init__() self.num_embeddings = num_embeddings self.max_relative_feature = max_relative_feature self.linear = hk.Linear(num_embeddings, name='embedding_linear') def __call__(self, offset, mask): d = jnp.clip(offset + self.max_relative_feature, 0, 2*self.max_relative_feature) * mask + \ (1 - mask) * (2*self.max_relative_feature + 1) d_onehot = jax.nn.one_hot(d, 2*self.max_relative_feature + 1 + 1) E = self.linear(d_onehot) return E class RunModel: def __init__(self, config) -> None: self.config = config def _forward_score(inputs): model = ProteinMPNN(**self.config) return model(**inputs) self.score = jax.jit(hk.transform(_forward_score).apply) self.init_score = jax.jit(hk.transform(_forward_score).init) def _forward_sample(inputs): model = ProteinMPNN(**self.config) return model.sample(**inputs) self.sample = jax.jit(hk.transform(_forward_sample).apply) self.init_sample = jax.jit(hk.transform(_forward_sample).init) def _forward_tsample(inputs): model = ProteinMPNN(**self.config) return model.tied_sample(**inputs) self.tied_sample = jax.jit(hk.transform(_forward_tsample).apply) self.init_tsample = jax.jit(hk.transform(_forward_tsample).init) def load_params(self, path): self.params = joblib.load(path) class ProteinFeatures(hk.Module): def __init__(self, edge_features, node_features, num_positional_embeddings=16, num_rbf=16, top_k=30, augment_eps=0., num_chain_embeddings=16): """ Extract protein features """ super(ProteinFeatures, self).__init__() self.edge_features = edge_features self.node_features = node_features self.top_k = top_k self.augment_eps = augment_eps self.num_rbf = num_rbf self.num_positional_embeddings = num_positional_embeddings self.embeddings = PositionalEncodings(num_positional_embeddings) node_in, edge_in = 6, num_positional_embeddings + num_rbf*25 self.edge_embedding = hk.Linear(edge_features, with_bias=False, name='edge_embedding') self.norm_edges = hk.LayerNorm(-1, create_scale=True, create_offset=True, name='norm_edges') self.safe_key = SafeKey(hk.next_rng_key()) def _get_edge_idx(self, X, mask, eps=1E-6): ''' get edge index input: mask.shape = (...,L), X.shape = (...,L,3) return: (...,L,k) ''' mask_2D = mask[...,None,:] * mask[...,:,None] dX = X[...,None,:,:] - X[...,:,None,:] D = jnp.sqrt(jnp.square(dX).sum(-1) + eps) D_masked = jnp.where(mask_2D,D,D.max(-1,keepdims=True)) k = min(self.top_k, X.shape[-2]) _, E_idx = jax.lax.approx_min_k(D_masked, k, reduction_dimension=-1) return E_idx def _rbf(self, D): ''' radial basis function (RBF) input: (...,L,k) output: (...,L,k,?) ''' D_min, D_max, D_count = 2., 22., self.num_rbf D_mu = jnp.linspace(D_min, D_max, D_count) D_sigma = (D_max - D_min) / D_count return jnp.exp(-((D[...,None] - D_mu) / D_sigma)**2) def _get_rbf(self, A, B, E_idx): D = jnp.sqrt(jnp.square(A[...,:,None,:] - B[...,None,:,:]).sum(-1) + 1e-6) D_neighbors = gather_edges(D[...,None], E_idx)[...,0] #[...,L,K] return self._rbf(D_neighbors) def __call__(self, X, mask, residue_idx, chain_idx, offset=None): if self.augment_eps > 0: self.safe_key, use_key = self.safe_key.split() X = X + self.augment_eps * jax.random.normal(use_key, X.shape) ########################## # get atoms ########################## # N,Ca,C,O,Cb Y = X.transpose((2,0,1,3)) if Y.shape[0] == 4: # add Cb b,c = (Y[1]-Y[0]),(Y[2]-Y[1]) Cb = -0.58273431*jnp.cross(b,c) + 0.56802827*b - 0.54067466*c + Y[1] Y = jnp.concatenate([Y,Cb[None]],0) ########################## # gather edge features ########################## # get edge indices (based on ca-ca distances) E_idx = self._get_edge_idx(Y[1], mask) # rbf encode distances between atoms edges = jnp.array([[1,1],[0,0],[2,2],[3,3],[4,4], [1,0],[1,2],[1,3],[1,4],[0,2], [0,3],[0,4],[4,2],[4,3],[3,2], [0,1],[2,1],[3,1],[4,1],[2,0], [3,0],[4,0],[2,4],[3,4],[2,3]]) RBF_all = jax.vmap(lambda x:self._get_rbf(Y[x[0]],Y[x[1]],E_idx))(edges) RBF_all = RBF_all.transpose((1,2,3,0,4)) RBF_all = RBF_all.reshape(RBF_all.shape[:-2]+(-1,)) ########################## # position embedding ########################## # residue index offset if offset is None: offset = (residue_idx[...,:,None] - residue_idx[...,None,:]) offset = gather_edges(offset[...,None], E_idx)[...,0] #[B, L, K] # chain index offset d_chains = (chain_idx[...,:,None] == chain_idx[...,None,:]).astype(int) E_chains = gather_edges(d_chains[...,None], E_idx)[...,0] E_positional = self.embeddings(offset, E_chains) ########################## # define edges ########################## E = jnp.concatenate((E_positional, RBF_all), -1) E = self.edge_embedding(E) E = self.norm_edges(E) return E, E_idx class EmbedToken(hk.Module): def __init__(self, vocab_size, embed_dim): super().__init__() self.vocab_size = vocab_size self.embed_dim = embed_dim self.w_init = hk.initializers.TruncatedNormal() @property def embeddings(self): return hk.get_parameter("W_s", [self.vocab_size, self.embed_dim], init=self.w_init) def __call__(self, arr): if jnp.issubdtype(arr.dtype, jnp.integer): one_hot = jax.nn.one_hot(arr, self.vocab_size) else: one_hot = arr return jnp.tensordot(one_hot, self.embeddings, 1) class ProteinMPNN(hk.Module, mpnn_sample): def __init__(self, num_letters, node_features, edge_features, hidden_dim, num_encoder_layers=3, num_decoder_layers=3, vocab=21, k_neighbors=64, augment_eps=0.05, dropout=0.1): super(ProteinMPNN, self).__init__() # Hyperparameters self.node_features = node_features self.edge_features = edge_features self.hidden_dim = hidden_dim # Featurization layers self.features = ProteinFeatures(edge_features, node_features, top_k=k_neighbors, augment_eps=augment_eps) self.W_e = hk.Linear(hidden_dim, with_bias=True, name='W_e') self.W_s = EmbedToken(vocab_size=vocab, embed_dim=hidden_dim) # Encoder layers self.encoder_layers = [ EncLayer(hidden_dim, hidden_dim*2, dropout=dropout, name='enc' + str(i)) for i in range(num_encoder_layers) ] # Decoder layers self.decoder_layers = [ DecLayer(hidden_dim, hidden_dim*3, dropout=dropout, name='dec' + str(i)) for i in range(num_decoder_layers) ] self.W_out = hk.Linear(num_letters, with_bias=True, name='W_out') def __call__(self, X, mask, residue_idx, chain_idx, S=None, chain_M=None, randn=None, ar_mask=None, decoding_order=None, offset=None): """ Graph-conditioned sequence model """ # Prepare node and edge embeddings E, E_idx = self.features(X, mask, residue_idx, chain_idx, offset=offset) h_V = jnp.zeros((E.shape[0], E.shape[1], E.shape[-1])) h_E = self.W_e(E) # Encoder is unmasked self-attention mask_attend = gather_nodes(mask[...,None],E_idx)[...,0] mask_attend = mask[...,None] * mask_attend for layer in self.encoder_layers: h_V, h_E = layer(h_V, h_E, E_idx, mask, mask_attend) # Build encoder embeddings h_EX_encoder = cat_neighbors_nodes(jnp.zeros_like(h_V), h_E, E_idx) h_EXV_encoder = cat_neighbors_nodes(h_V, h_EX_encoder, E_idx) if S is None: ########################################## # unconditional_probs ########################################## # make an autogressive mask ar_mask = jnp.zeros([X.shape[0], X.shape[1], X.shape[1]]) mask_attend = jnp.take_along_axis(ar_mask, E_idx, 2)[...,None] mask_1D = mask.reshape([mask.shape[0], mask.shape[1], 1, 1]) mask_bw = mask_1D * mask_attend mask_fw = mask_1D * (1. - mask_attend) h_EXV_encoder_fw = mask_fw * h_EXV_encoder for layer in self.decoder_layers: h_V = layer(h_V, h_EXV_encoder_fw, mask) else: ########################################## # conditional_probs ########################################## # Concatenate sequence embeddings for autoregressive decoder h_S = self.W_s(S) h_ES = cat_neighbors_nodes(h_S, h_E, E_idx) if ar_mask is None: if decoding_order is None: # update chain_M to include missing regions chain_M = chain_M * mask #[numbers will be smaller for places where chain_M = 0.0 and higher for places where chain_M = 1.0] decoding_order = jnp.argsort((chain_M+0.0001)*(jnp.abs(randn))) # make an autogressive mask ar_mask = get_ar_mask(decoding_order) mask_attend = jnp.take_along_axis(ar_mask, E_idx, 2)[...,None] mask_1D = mask.reshape([mask.shape[0], mask.shape[1], 1, 1]) mask_bw = mask_1D * mask_attend mask_fw = mask_1D * (1. - mask_attend) h_EXV_encoder_fw = mask_fw * h_EXV_encoder for layer in self.decoder_layers: # Masked positions attend to encoder information, unmasked see. h_ESV = cat_neighbors_nodes(h_V, h_ES, E_idx) h_ESV = mask_bw * h_ESV + h_EXV_encoder_fw h_V = layer(h_V, h_ESV, mask) logits = self.W_out(h_V) log_probs = jax.nn.log_softmax(logits, axis=-1) return logits, log_probs