def net_rbf(net_p, net_q, J=5):
'''
Network equivalent to Wittawat's mean embedding test:
compute RBF kernel values to each of J test points.
'''
from layers import RBFLayer
net_p = RBFLayer(net_p, J)
net_q = RBFLayer(net_q, J, locs=net_p.locs, log_sigma=net_p.log_sigma)
return net_p, net_q, 0
def __init__(self,
dim=DIM,
output_dim=OUTPUT_DIM,
edge_dim=EDGE_DIM,
cutoff=CUTOFF,
width=WIDTH,
n_conv=3,
norm=False,
atom_ref=None):
super().__init__()
self.name = "MGCN"
self._dim, self.output_dim = dim, output_dim
self.edge_dim, self.cutoff, self.atom_ref = edge_dim, cutoff, atom_ref
self.width = width
self.n_conv, self.norm = n_conv, norm
self.activation = nn.Softplus(beta=1, threshold=THRESHOLD)
if atom_ref is not None:
self.e0 = AtomEmbedding(1)
self.embedding_layer, self.edge_embedding_layer = AtomEmbedding(
dim), EdgeEmbedding(dim=edge_dim)
self.rbf_layer = RBFLayer(0, cutoff, width)
self.conv_layers = nn.ModuleList([
MultiLevelInteraction(self.rbf_layer._fan_out, dim)
for i in range(n_conv)
])
self.node_dense_layer1, self.node_dense_layer2 = nn.Linear(
dim * (self.n_conv + 1), NDL), nn.Linear(NDL, output_dim)
def __init__(self,
dim=128,
output_dim=1,
edge_dim=128,
cutoff=5.0,
width=1,
n_conv=3,
norm=False,
atom_ref=None,
pre_train=None):
"""
Args:
dim: dimension of feature maps
cutoff: the edge larger than cutoff will be delete
width: width in the RBF layer
atom_ref: atom reference
used as the initial value of atom embeddings,
Or set to None with random initialization
norm: normalization
"""
super().__init__()
self.name = "MGCN"
self._dim = dim
self.output_dim = output_dim
self.edge_dim = edge_dim
self.cutoff = cutoff
self.width = width
self.n_conv = n_conv
self.atom_ref = atom_ref
self.norm = norm
self.register_hyper_params(dim=dim,
output_dim=output_dim,
edge_dim=edge_dim,
cutoff=cutoff,
width=width,
n_conv=n_conv,
atom_ref=atom_ref,
norm=norm)
self.activation = nn.Softplus(beta=1, threshold=20)
if atom_ref is not None:
self.e0 = AtomEmbedding(1, pre_train=atom_ref)
if pre_train is None:
self.embedding_layer = AtomEmbedding(dim)
else:
self.embedding_layer = AtomEmbedding(pre_train=pre_train)
self.edge_embedding_layer = EdgeEmbedding(dim=edge_dim)
self.rbf_layer = RBFLayer(0, cutoff, width)
self.conv_layers = nn.ModuleList([
MultiLevelInteraction(self.rbf_layer._fan_out, dim)
for i in range(n_conv)
])
self.node_dense_layer1 = nn.Linear(dim * (self.n_conv + 1), 64)
self.node_dense_layer2 = nn.Linear(64, output_dim)
def __init__(self,
dim=128,
output_dim=1,
edge_dim=128,
cutoff=5.0,
width=1,
n_conv=3,
norm=False,
atom_ref=None,
pre_train=None):
"""
Args:
dim: dimension of feature maps
out_put_dim: the num of target propperties to predict
edge_dim: dimension of edge feature
cutoff: the maximum distance between nodes
width: width in the RBF layer
n_conv: number of convolutional layers
norm: normalization
atom_ref: atom reference
used as the initial value of atom embeddings,
or set to None with random initialization
pre_train: pre_trained node embeddings
"""
super().__init__()
self.name = "MGCN"
self._dim = dim
self.output_dim = output_dim
self.edge_dim = edge_dim
self.cutoff = cutoff
self.width = width
self.n_conv = n_conv
self.atom_ref = atom_ref
self.norm = norm
self.activation = nn.Softplus(beta=1, threshold=20)
if atom_ref is not None:
self.e0 = AtomEmbedding(1, pre_train=atom_ref)
if pre_train is None:
self.embedding_layer = AtomEmbedding(dim)
else:
self.embedding_layer = AtomEmbedding(pre_train=pre_train)
self.edge_embedding_layer = EdgeEmbedding(dim=edge_dim)
self.rbf_layer = RBFLayer(0, cutoff, width)
self.conv_layers = nn.ModuleList([
MultiLevelInteraction(self.rbf_layer._fan_out, dim)
for i in range(n_conv)
])
self.node_dense_layer1 = nn.Linear(dim * (self.n_conv + 1), 64)
self.node_dense_layer2 = nn.Linear(64, output_dim)
def __init__(self,
dim=64,
cutoff=5.0,
output_dim=1,
width=1,
n_conv=3,
norm=False,
atom_ref=None,
pre_train=None):
"""
Args:
dim: dimension of features
output_dim: dimension of prediction
cutoff: radius cutoff
width: width in the RBF function
n_conv: number of interaction layers
atom_ref: used as the initial value of atom embeddings,
or set to None with random initialization
norm: normalization
"""
super().__init__()
self.name = "SchNet"
self._dim = dim
self.cutoff = cutoff
self.width = width
self.n_conv = n_conv
self.atom_ref = atom_ref
self.norm = norm
self.activation = ShiftSoftplus()
if atom_ref is not None:
self.e0 = AtomEmbedding(1, pre_train=atom_ref)
if pre_train is None:
self.embedding_layer = AtomEmbedding(dim)
else:
self.embedding_layer = AtomEmbedding(pre_train=pre_train)
self.rbf_layer = RBFLayer(0, cutoff, width)
self.conv_layers = nn.ModuleList(
[Interaction(self.rbf_layer._fan_out, dim) for i in range(n_conv)])
self.atom_dense_layer1 = nn.Linear(dim, 64)
self.atom_dense_layer2 = nn.Linear(64, output_dim)
self.fc = nn.Sequential(
nn.Linear(69, 128),
nn.ReLU(),
nn.Linear(128, 64),
nn.ReLU(),
nn.Dropout(0.2),
nn.Linear(64, 32),
nn.ReLU(),
nn.Linear(32, 1),
)