File size: 7,756 Bytes
4742cab |
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 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 |
import torch
import torch.nn as nn
import torch.nn.functional as F
from equivariant_diffusion.egnn_new import EGNN, GNN
from equivariant_diffusion.en_diffusion import EnVariationalDiffusion
remove_mean_batch = EnVariationalDiffusion.remove_mean_batch
import numpy as np
class EGNNDynamics(nn.Module):
def __init__(self, atom_nf, residue_nf,
n_dims, joint_nf=16, hidden_nf=64, device='cpu',
act_fn=torch.nn.SiLU(), n_layers=4, attention=False,
condition_time=True, tanh=False, mode='egnn_dynamics',
norm_constant=0, inv_sublayers=2, sin_embedding=False,
normalization_factor=100, aggregation_method='sum',
update_pocket_coords=True, edge_cutoff_ligand=None,
edge_cutoff_pocket=None, edge_cutoff_interaction=None,
reflection_equivariant=True, edge_embedding_dim=None):
super().__init__()
self.mode = mode
self.edge_cutoff_l = edge_cutoff_ligand
self.edge_cutoff_p = edge_cutoff_pocket
self.edge_cutoff_i = edge_cutoff_interaction
self.edge_nf = edge_embedding_dim
self.atom_encoder = nn.Sequential(
nn.Linear(atom_nf, 2 * atom_nf),
act_fn,
nn.Linear(2 * atom_nf, joint_nf)
)
self.atom_decoder = nn.Sequential(
nn.Linear(joint_nf, 2 * atom_nf),
act_fn,
nn.Linear(2 * atom_nf, atom_nf)
)
self.residue_encoder = nn.Sequential(
nn.Linear(residue_nf, 2 * residue_nf),
act_fn,
nn.Linear(2 * residue_nf, joint_nf)
)
self.residue_decoder = nn.Sequential(
nn.Linear(joint_nf, 2 * residue_nf),
act_fn,
nn.Linear(2 * residue_nf, residue_nf)
)
self.edge_embedding = nn.Embedding(3, self.edge_nf) \
if self.edge_nf is not None else None
self.edge_nf = 0 if self.edge_nf is None else self.edge_nf
if condition_time:
dynamics_node_nf = joint_nf + 1
else:
print('Warning: dynamics model is _not_ conditioned on time.')
dynamics_node_nf = joint_nf
if mode == 'egnn_dynamics':
self.egnn = EGNN(
in_node_nf=dynamics_node_nf, in_edge_nf=self.edge_nf,
hidden_nf=hidden_nf, device=device, act_fn=act_fn,
n_layers=n_layers, attention=attention, tanh=tanh,
norm_constant=norm_constant,
inv_sublayers=inv_sublayers, sin_embedding=sin_embedding,
normalization_factor=normalization_factor,
aggregation_method=aggregation_method,
reflection_equiv=reflection_equivariant
)
self.node_nf = dynamics_node_nf
self.update_pocket_coords = update_pocket_coords
elif mode == 'gnn_dynamics':
self.gnn = GNN(
in_node_nf=dynamics_node_nf + n_dims, in_edge_nf=self.edge_nf,
hidden_nf=hidden_nf, out_node_nf=n_dims + dynamics_node_nf,
device=device, act_fn=act_fn, n_layers=n_layers,
attention=attention, normalization_factor=normalization_factor,
aggregation_method=aggregation_method)
self.device = device
self.n_dims = n_dims
self.condition_time = condition_time
def forward(self, xh_atoms, xh_residues, t, mask_atoms, mask_residues):
x_atoms = xh_atoms[:, :self.n_dims].clone()
h_atoms = xh_atoms[:, self.n_dims:].clone()
x_residues = xh_residues[:, :self.n_dims].clone()
h_residues = xh_residues[:, self.n_dims:].clone()
# embed atom features and residue features in a shared space
h_atoms = self.atom_encoder(h_atoms)
h_residues = self.residue_encoder(h_residues)
# combine the two node types
x = torch.cat((x_atoms, x_residues), dim=0)
h = torch.cat((h_atoms, h_residues), dim=0)
mask = torch.cat([mask_atoms, mask_residues])
if self.condition_time:
if np.prod(t.size()) == 1:
# t is the same for all elements in batch.
h_time = torch.empty_like(h[:, 0:1]).fill_(t.item())
else:
# t is different over the batch dimension.
h_time = t[mask]
h = torch.cat([h, h_time], dim=1)
# get edges of a complete graph
edges = self.get_edges(mask_atoms, mask_residues, x_atoms, x_residues)
assert torch.all(mask[edges[0]] == mask[edges[1]])
# Get edge types
if self.edge_nf > 0:
# 0: ligand-pocket, 1: ligand-ligand, 2: pocket-pocket
edge_types = torch.zeros(edges.size(1), dtype=int, device=edges.device)
edge_types[(edges[0] < len(mask_atoms)) & (edges[1] < len(mask_atoms))] = 1
edge_types[(edges[0] >= len(mask_atoms)) & (edges[1] >= len(mask_atoms))] = 2
# Learnable embedding
edge_types = self.edge_embedding(edge_types)
else:
edge_types = None
if self.mode == 'egnn_dynamics':
update_coords_mask = None if self.update_pocket_coords \
else torch.cat((torch.ones_like(mask_atoms),
torch.zeros_like(mask_residues))).unsqueeze(1)
h_final, x_final = self.egnn(h, x, edges,
update_coords_mask=update_coords_mask,
batch_mask=mask, edge_attr=edge_types)
vel = (x_final - x)
elif self.mode == 'gnn_dynamics':
xh = torch.cat([x, h], dim=1)
output = self.gnn(xh, edges, node_mask=None, edge_attr=edge_types)
vel = output[:, :3]
h_final = output[:, 3:]
else:
raise Exception("Wrong mode %s" % self.mode)
if self.condition_time:
# Slice off last dimension which represented time.
h_final = h_final[:, :-1]
# decode atom and residue features
h_final_atoms = self.atom_decoder(h_final[:len(mask_atoms)])
h_final_residues = self.residue_decoder(h_final[len(mask_atoms):])
if torch.any(torch.isnan(vel)):
if self.training:
vel[torch.isnan(vel)] = 0.0
else:
raise ValueError("NaN detected in EGNN output")
if self.update_pocket_coords:
# in case of unconditional joint distribution, include this as in
# the original code
vel = remove_mean_batch(vel, mask)
return torch.cat([vel[:len(mask_atoms)], h_final_atoms], dim=-1), \
torch.cat([vel[len(mask_atoms):], h_final_residues], dim=-1)
def get_edges(self, batch_mask_ligand, batch_mask_pocket, x_ligand, x_pocket):
adj_ligand = batch_mask_ligand[:, None] == batch_mask_ligand[None, :]
adj_pocket = batch_mask_pocket[:, None] == batch_mask_pocket[None, :]
adj_cross = batch_mask_ligand[:, None] == batch_mask_pocket[None, :]
if self.edge_cutoff_l is not None:
adj_ligand = adj_ligand & (torch.cdist(x_ligand, x_ligand) <= self.edge_cutoff_l)
if self.edge_cutoff_p is not None:
adj_pocket = adj_pocket & (torch.cdist(x_pocket, x_pocket) <= self.edge_cutoff_p)
if self.edge_cutoff_i is not None:
adj_cross = adj_cross & (torch.cdist(x_ligand, x_pocket) <= self.edge_cutoff_i)
adj = torch.cat((torch.cat((adj_ligand, adj_cross), dim=1),
torch.cat((adj_cross.T, adj_pocket), dim=1)), dim=0)
edges = torch.stack(torch.where(adj), dim=0)
return edges
|