change num_edge_type to be argument
@@ -24,13 +24,15 @@ class GGNN(chainer.Chain):
concat_hidden (bool): If set to True, readout is executed in each layer
and the result is concatenated
weight_tying (bool): enable weight_tying or not
+ num_edge_type (int): number of edge type.
+ Defaults to 4 for single, double, triple and aromatic bond.
"""
- NUM_EDGE_TYPE = 4
def __init__(self, out_dim, hidden_dim=16, n_layers=4,
n_atom_types=MAX_ATOMIC_NUM, concat_hidden=False,
- weight_tying=True, activation=functions.identity):
+ weight_tying=True, activation=functions.identity,
+ num_edge_type=4):
super(GGNN, self).__init__()
n_readout_layer = n_layers if concat_hidden else 1
n_message_layer = 1 if weight_tying else n_layers
@@ -38,7 +40,7 @@ class GGNN(chainer.Chain):
# Update
self.embed = EmbedAtomID(out_size=hidden_dim, in_size=n_atom_types)
self.update_layers = chainer.ChainList(*[GGNNUpdate(
- hidden_dim=hidden_dim, num_edge_type=self.NUM_EDGE_TYPE)
+ hidden_dim=hidden_dim, num_edge_type=num_edge_type)
for _ in range(n_message_layer)])
# Readout
self.readout_layers = chainer.ChainList(*[GGNNReadout(
@@ -48,6 +50,7 @@ class GGNN(chainer.Chain):
self.out_dim = out_dim
self.hidden_dim = hidden_dim
self.n_layers = n_layers
+ self.num_edge_type = num_edge_type
self.activation = activation
self.concat_hidden = concat_hidden
self.weight_tying = weight_tying
