use GraphConvModel for schnet
-import chainer
-from chainer import functions
-
from chainer_chemistry.config import MAX_ATOMIC_NUM
-from chainer_chemistry.links.connection.embed_atom_id import EmbedAtomID
from chainer_chemistry.links.readout.schnet_readout import SchNetReadout
from chainer_chemistry.links.update.schnet_update import SchNetUpdate
+from chainer_chemistry.models.graph_conv_model import GraphConvModel
-class SchNet(chainer.Chain):
+class SchNet(GraphConvModel):
"""SchNet
See Kristof et al, \
@@ -32,35 +29,15 @@ class SchNet(chainer.Chain):
gamma (float): exponential factor of `CFConv`'s radius kernel.
"""
- def __init__(self, out_dim=1, hidden_dim=64, n_layers=3,
+ def __init__(self, out_dim=1, in_channels=64, n_layers=3,
readout_hidden_dim=32, n_atom_types=MAX_ATOMIC_NUM,
concat_hidden=False, num_rbf=300, radius_resolution=0.1,
- gamma=10.0):
- super(SchNet, self).__init__()
- with self.init_scope():
- self.embed = EmbedAtomID(out_size=hidden_dim, in_size=n_atom_types)
- self.update_layers = chainer.ChainList(
- *[SchNetUpdate(
- hidden_dim,
- num_rbf=num_rbf, radius_resolution=radius_resolution,
- gamma=gamma) for _ in range(n_layers)])
- self.readout_layer = SchNetReadout(out_dim, readout_hidden_dim)
- self.out_dim = out_dim
- self.hidden_dim = hidden_dim
- self.readout_hidden_dim = readout_hidden_dim
- self.n_layers = n_layers
- self.concat_hidden = concat_hidden
-
- def __call__(self, atom_features, dist_features):
- x = self.embed(atom_features)
- h = []
- # --- update part ---
- for i in range(self.n_layers):
- x = self.update_layers[i](x, dist_features)
- if self.concat_hidden:
- h.append(x)
- # --- readout part ---
- if self.concat_hidden:
- x = functions.concat(h, axis=2)
- x = self.readout_layer(x)
- return x
+ gamma=10.0, with_gwm=False):
+ # TODO: use readout_hidden_dim
+ # TODO: use num_rbf, radius_resolution, gamma
+ super(SchNet, self).__init__(
+ update_layer=SchNetUpdate, readout_layer=SchNetReadout,
+ out_dim=out_dim, in_channels=in_channels, n_layers=n_layers,
+ n_atom_types=n_atom_types, concat_hidden=concat_hidden,
+ with_gwm=with_gwm
+ )
