File size: 2,223 Bytes
39c21b2 | 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 | import torch
import torch.nn as nn
import torch.nn.functional as F
from cdvae.pl_modules.embeddings import MAX_ATOMIC_NUM
from cdvae.pl_modules.gemnet.gemnet import GemNetT
def build_mlp(in_dim, hidden_dim, fc_num_layers, out_dim):
mods = [nn.Linear(in_dim, hidden_dim), nn.ReLU()]
for i in range(fc_num_layers-1):
mods += [nn.Linear(hidden_dim, hidden_dim), nn.ReLU()]
mods += [nn.Linear(hidden_dim, out_dim)]
return nn.Sequential(*mods)
class GemNetTDecoder(nn.Module):
"""Decoder with GemNetT."""
def __init__(
self,
hidden_dim=128,
latent_dim=256,
max_neighbors=20,
radius=6.,
scale_file=None,
):
super(GemNetTDecoder, self).__init__()
self.cutoff = radius
self.max_num_neighbors = max_neighbors
self.gemnet = GemNetT(
num_targets=1,
latent_dim=latent_dim,
emb_size_atom=hidden_dim,
emb_size_edge=hidden_dim,
regress_forces=True,
cutoff=self.cutoff,
max_neighbors=self.max_num_neighbors,
otf_graph=True,
scale_file=scale_file,
)
self.fc_atom = nn.Linear(hidden_dim, MAX_ATOMIC_NUM)
def forward(self, z, pred_frac_coords, pred_atom_types, num_atoms,
lengths, angles):
"""
args:
z: (N_cryst, num_latent)
pred_frac_coords: (N_atoms, 3)
pred_atom_types: (N_atoms, ), need to use atomic number e.g. H = 1
num_atoms: (N_cryst,)
lengths: (N_cryst, 3)
angles: (N_cryst, 3)
returns:
atom_frac_coords: (N_atoms, 3)
atom_types: (N_atoms, MAX_ATOMIC_NUM)
"""
# (num_atoms, hidden_dim) (num_crysts, 3)
h, pred_cart_coord_diff = self.gemnet(
z=z,
frac_coords=pred_frac_coords,
atom_types=pred_atom_types,
num_atoms=num_atoms,
lengths=lengths,
angles=angles,
edge_index=None,
to_jimages=None,
num_bonds=None,
)
pred_atom_types = self.fc_atom(h)
return pred_cart_coord_diff, pred_atom_types
|