| import jax
|
| import jax.numpy as jnp
|
| import haiku as hk
|
| import numpy as np
|
|
|
| from .utils import cat_neighbors_nodes, get_ar_mask
|
|
|
| class mpnn_sample:
|
| def sample(self, I):
|
| """
|
| I = {
|
| [[required]]
|
| 'X' = (L,4,3)
|
| 'mask' = (L,)
|
| 'residue_index' = (L,)
|
| 'chain_idx' = (L,)
|
| 'decoding_order' = (L,)
|
|
|
| [[optional]]
|
| 'ar_mask' = (L,L)
|
| 'bias' = (L,21)
|
| 'temperature' = 1.0
|
| }
|
| """
|
|
|
| key = hk.next_rng_key()
|
| L = I["X"].shape[0]
|
| temperature = I.get("temperature",1.0)
|
|
|
|
|
| E, E_idx = self.features(I)
|
| h_V = jnp.zeros((E.shape[0], E.shape[-1]))
|
| h_E = self.W_e(E)
|
|
|
|
|
|
|
|
|
| mask_attend = jnp.take_along_axis(I["mask"][:,None] * I["mask"][None,:], E_idx, 1)
|
| for layer in self.encoder_layers:
|
| h_V, h_E = layer(h_V, h_E, E_idx, I["mask"], mask_attend)
|
|
|
|
|
| ar_mask = I.get("ar_mask",get_ar_mask(I["decoding_order"]))
|
|
|
| mask_attend = jnp.take_along_axis(ar_mask, E_idx, 1)
|
| mask_1D = I["mask"][:,None]
|
| mask_bw = mask_1D * mask_attend
|
| mask_fw = mask_1D * (1 - mask_attend)
|
|
|
| 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)
|
| h_EXV_encoder = mask_fw[...,None] * h_EXV_encoder
|
|
|
| def fwd(x, t, key):
|
| h_EXV_encoder_t = h_EXV_encoder[t]
|
| E_idx_t = E_idx[t]
|
| mask_t = I["mask"][t]
|
| mask_bw_t = mask_bw[t]
|
| h_ES_t = cat_neighbors_nodes(x["h_S"], h_E[t], E_idx_t)
|
|
|
|
|
|
|
|
|
| for l,layer in enumerate(self.decoder_layers):
|
| h_V = x["h_V"][l]
|
| h_ESV_decoder_t = cat_neighbors_nodes(h_V, h_ES_t, E_idx_t)
|
| h_ESV_t = mask_bw_t[...,None] * h_ESV_decoder_t + h_EXV_encoder_t
|
| h_V_t = layer(h_V[t], h_ESV_t, mask_V=mask_t)
|
|
|
| x["h_V"] = x["h_V"].at[l+1,t].set(h_V_t)
|
|
|
| logits_t = self.W_out(h_V_t)
|
| x["logits"] = x["logits"].at[t].set(logits_t)
|
|
|
|
|
|
|
|
|
|
|
|
|
| if "bias" in I: logits_t += I["bias"][t]
|
|
|
|
|
| logits_t = logits_t/temperature + jax.random.gumbel(key, logits_t.shape)
|
|
|
|
|
| logits_t = logits_t.mean(0, keepdims=True)
|
|
|
| S_t = jax.nn.one_hot(logits_t[...,:20].argmax(-1), 21)
|
|
|
|
|
| x["h_S"] = x["h_S"].at[t].set(self.W_s(S_t))
|
| x["S"] = x["S"].at[t].set(S_t)
|
| return x, None
|
|
|
|
|
| X = {"h_S": jnp.zeros_like(h_V),
|
| "h_V": jnp.array([h_V] + [jnp.zeros_like(h_V)] * len(self.decoder_layers)),
|
| "S": jnp.zeros((L,21)),
|
| "logits": jnp.zeros((L,21))}
|
|
|
|
|
| t = I["decoding_order"]
|
| if t.ndim == 1: t = t[:,None]
|
| XS = {"t":t, "key":jax.random.split(key,t.shape[0])}
|
| X = hk.scan(lambda x, xs: fwd(x, xs["t"], xs["key"]), X, XS)[0]
|
|
|
| return {"S":X["S"], "logits":X["logits"], "decoding_order":t} |
