File size: 4,330 Bytes
8e04e6f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 | from typing import Dict
import einops
import torch
from onescience.utils.openfold.np.residue_constants import RESTYPE_ATOM37_MASK
from models.nn.feature_factory import FeatureFactory
def get_atom_mask(device: torch.device = None):
return torch.from_numpy(RESTYPE_ATOM37_MASK).to(
dtype=torch.bool, device=device
) # [21, 37]
class ResidualLayer(torch.nn.Module):
"""Residual layer with pre-LN"""
def __init__(self, dim):
super(ResidualLayer, self).__init__()
self.layer_norm = torch.nn.LayerNorm(dim)
self.linear = torch.nn.Linear(dim, dim)
self.softplus = torch.nn.Softplus()
def forward(self, x):
out = self.layer_norm(x)
out = self.linear(out)
out = self.softplus(out)
return x + out
class DecoderFFLocal(torch.nn.Module):
"""
Encoder part of the autoencoder. A transformer with pair-biased attention.
"""
def __init__(self, **kwargs):
"""
Initializes the NN. The seqs and pair representations used are just zero in case
no features are required."""
super(DecoderFFLocal, self).__init__()
nlayers = kwargs["decoder"]["nlayers"]
token_dim = kwargs["decoder"]["token_dim"] # hidden dim
# To form initial representation
self.init_repr_factory = FeatureFactory(
feats=kwargs["decoder"]["feats_seq"],
dim_feats_out=token_dim,
use_ln_out=False,
mode="seq",
**kwargs["decoder"],
)
layers = []
for _ in range(nlayers):
layers.append(ResidualLayer(dim=token_dim))
self.ff_nn = torch.nn.Sequential(*layers)
self.logit_linear = torch.nn.Sequential(
torch.nn.LayerNorm(token_dim),
torch.nn.Linear(token_dim, 20, bias=False),
)
self.struct_linear = torch.nn.Sequential(
torch.nn.LayerNorm(token_dim),
torch.nn.Linear(token_dim, int(37 * 3), bias=False),
)
def forward(self, input: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
"""
Runs the network.
Args:
input: {
"z_latent": torch.Tensor(b, n, latent_dim),
"ca_coors_nm": torch.Tensor(b, n, 3),
"residue_mask": boolean torch.Tensor(b, n)
}
Returns:
Dictionary:
{
"coors_nm": all atom coordinates, shape [b, n, 37, 3]
"seq_logits": logits for the residue types, shape [b, n, 20]
"residue_mask": boolean [b, n]
"aatype_max": residue type by taking the most likely logit, shape [b, n], with integer values {0, ..., 19}
"atom_mask": boolean [b, n, 37], atom37 mask corresponding to aatype_max
}
"""
ca_coors_nm = input["ca_coors_nm"] # [b, n, 3]
mask = input["residue_mask"] # [b, n] boolean
# Iinitial sequence representation from features
seq_f_repr = self.init_repr_factory(input) # [b, n, token_dim]
seqs = seq_f_repr * mask[..., None] # [b, n, token_dim]
seqs = self.ff_nn(seqs) * mask[..., None] # [b, n, token_dim]
# Get logits
logits_out = self.logit_linear(seqs) * mask[..., None] # [b, n, 20]
# Get coordinates
coors_flat_nm = self.struct_linear(seqs) * mask[..., None] # [b, n, 37 * 3]
coors_a37_nm = einops.rearrange(
coors_flat_nm, "b n (a t) -> b n a t", a=37, t=3
) # [b, n, 37, 3]
coors_a37_nm[..., 1, :] = coors_a37_nm[..., 1, :] * 0.0 + ca_coors_nm
# Get sequence
aatype_max = torch.argmax(logits_out, dim=-1) # [b, n]
aatype_max = aatype_max * mask # [b, n]
# Get atom_mask
aa_a37_mask = get_atom_mask(device=logits_out.device) # [21, 37] boolean
atom_mask = aa_a37_mask[aatype_max, :] # [b, n, 37] boolean
atom_mask = atom_mask * mask[..., None] # [b, n, 37] boolean
output = {
"coors_nm": coors_a37_nm, # [b, n, 37, 3]
"seq_logits": logits_out, # [b, n, 20]
"residue_mask": mask, # [b, n]
"aatype_max": aatype_max, # [b, n]
"atom_mask": atom_mask, # [b, n, 37]
}
return output
|