wehub-resource-sync
61 lines
1.8 KiB
Python
61 lines
1.8 KiB
Python
import numpy as np
|
|
import torch
|
|
import torch.nn as nn
|
|
from modules.envelope import Envelope
|
|
from modules.initializers import GlorotOrthogonal
|
|
|
|
|
|
class EmbeddingBlock(nn.Module):
|
|
def __init__(
|
|
self,
|
|
emb_size,
|
|
num_radial,
|
|
bessel_funcs,
|
|
cutoff,
|
|
envelope_exponent,
|
|
num_atom_types=95,
|
|
activation=None,
|
|
):
|
|
super(EmbeddingBlock, self).__init__()
|
|
|
|
self.bessel_funcs = bessel_funcs
|
|
self.cutoff = cutoff
|
|
self.activation = activation
|
|
self.envelope = Envelope(envelope_exponent)
|
|
self.embedding = nn.Embedding(num_atom_types, emb_size)
|
|
self.dense_rbf = nn.Linear(num_radial, emb_size)
|
|
self.dense = nn.Linear(emb_size * 3, emb_size)
|
|
self.reset_params()
|
|
|
|
def reset_params(self):
|
|
nn.init.uniform_(self.embedding.weight, a=-np.sqrt(3), b=np.sqrt(3))
|
|
GlorotOrthogonal(self.dense_rbf.weight)
|
|
GlorotOrthogonal(self.dense.weight)
|
|
|
|
def edge_init(self, edges):
|
|
"""msg emb init"""
|
|
# m init
|
|
rbf = self.dense_rbf(edges.data["rbf"])
|
|
if self.activation is not None:
|
|
rbf = self.activation(rbf)
|
|
|
|
m = torch.cat([edges.src["h"], edges.dst["h"], rbf], dim=-1)
|
|
m = self.dense(m)
|
|
if self.activation is not None:
|
|
m = self.activation(m)
|
|
|
|
# rbf_env init
|
|
d_scaled = edges.data["d"] / self.cutoff
|
|
rbf_env = [f(d_scaled) for f in self.bessel_funcs]
|
|
rbf_env = torch.stack(rbf_env, dim=1)
|
|
|
|
d_cutoff = self.envelope(d_scaled)
|
|
rbf_env = d_cutoff[:, None] * rbf_env
|
|
|
|
return {"m": m, "rbf_env": rbf_env}
|
|
|
|
def forward(self, g):
|
|
g.ndata["h"] = self.embedding(g.ndata["Z"])
|
|
g.apply_edges(self.edge_init)
|
|
return g
|