def __init__(self, n_atom_basis, n_filters, n_gaussians, cutoff, trainable_gauss, ): super(SchNetConv, self).__init__() self.moduledict = ModuleDict({ 'message_edge_filter': Sequential( GaussianSmearing( start=0.0, stop=cutoff, n_gaussians=n_gaussians, trainable=trainable_gauss ), Dense(in_features=n_gaussians, out_features=n_gaussians), shifted_softplus(), Dense(in_features=n_gaussians, out_features=n_filters) ), 'message_node_filter': Dense(in_features=n_atom_basis, out_features=n_filters), 'update_function': Sequential( Dense(in_features=n_filters, out_features=n_atom_basis), shifted_softplus(), Dense(in_features=n_atom_basis, out_features=n_atom_basis) ) })
def __init__(self, cutoff, n_gaussians, trainable_gauss, n_filters, dropout_rate, activation='shifted_softplus'): super(SchNetEdgeFilter, self).__init__() self.filter = Sequential( GaussianSmearing( start=0.0, stop=cutoff, n_gaussians=n_gaussians, trainable=trainable_gauss, ), Dense( in_features=n_gaussians, out_features=n_gaussians, dropout_rate=dropout_rate, ), layer_types[activation](), Dense( in_features=n_gaussians, out_features=n_filters, dropout_rate=dropout_rate, ))
def __init__(self, n_atom_hidden, n_filters, dropout_rate, n_bond_hidden, activation='shifted_softplus'): """ Args: n_atom_hidden (int): hidden dimension of the atom features. Same as `n_atom_basis` in regular SchNet, but different than `n_atom_basis` in SchNetFeatures, where `n_atom_basis` is the initial dimension of the atom feature vector and `n_atom_hidden` is its dimension after going through another network. n_filters (int): dimension of the distance hidden vector dropout_rate (float): dropout rate n_bond_hidden (int): dimension of the bond hidden vector activation (str): nonlinear activation name Returns: None """ super(MixedSchNetConv, self).__init__() self.moduledict = ModuleDict({ # convert the atom features to the dimension # of cat(hidden_distance, hidden_bond) "message_node_filter": Dense( in_features=n_atom_hidden, out_features=(n_filters + n_bond_hidden), dropout_rate=dropout_rate, ), # after multiplying edge features with # node features, convert them back to size # `n_atom_hidden` "update_function": Sequential( Dense( in_features=(n_filters + n_bond_hidden), out_features=n_atom_hidden, dropout_rate=dropout_rate, ), layer_types[activation](), Dense( in_features=n_atom_hidden, out_features=n_atom_hidden, dropout_rate=dropout_rate, ), ), })
def __init__(self, n_edge_hidden, dropout_rate, activation, **kwargs): """ Args: n_edge_hidden: dimension of the hidden edge vector dropout_rate (float): dropout rate activation (str): name of non-linear activation function Returns: None """ MessagePassingModule.__init__(self) # As in the original ChemProp paper, # the features are added together linearly # with no bias. This means that features # equal to 0 don't contribute to the output. # This is important, for example, for # CpSchNetConv, in which every non- # bonded neighbour has zeros for its # ChemProp bond features. We don't want # these zeros contributing to the output. self.dense = Dense(in_features=n_edge_hidden, out_features=n_edge_hidden, dropout_rate=dropout_rate, bias=False) self.activation = layer_types[activation]()
def get_dense(inp_dim, out_dim, activation, bias): """ Create a dense layer. Args: inp_dim (int): dimension of input out_dim (int): dimension of output activation (str): name of activation layer bias (bool): whether or not to add a bias Returns: (nn.layers.Dense): dense layer """ if activation is not None: activation = layer_types[activation]() return Dense(inp_dim, out_dim, activation=activation, bias=bias)
def __init__(self, n_bond_hidden, cp_dropout, gauss_embed, cutoff, n_gaussians, trainable_gauss, n_filters, schnet_dropout, activation, **kwargs): """ Args: n_bond_hidden (int): bond feature hidden dimension cp_dropout (float): dropout rate for the ChemProp convolution gauss_embed (bool): whether to embed distances in a basis of Gaussians. cutoff (float): neighbor list cutoff n_gaussians (int): number of Gaussians in which to expand distances. trainable_gauss (bool): whether Gaussian spacings and widths are learnable parameters. n_filters (int): hidden distance feature dimension schnet_dropout (float): dropout rate for SchNet embedding activation (str): name of nonlinear activation function Returns: None """ ChemPropConv.__init__(self, n_edge_hidden=n_bond_hidden, dropout_rate=cp_dropout, activation=activation) self.n_bond_hidden = n_bond_hidden self.moduledict = ModuleDict({}) if not gauss_embed: return edge_filter = Sequential( GaussianSmearing( start=0.0, stop=cutoff, n_gaussians=n_gaussians, trainable=trainable_gauss, ), Dense( in_features=n_gaussians, out_features=n_filters, dropout_rate=schnet_dropout, ), layer_types[activation]()) self.moduledict["edge_filter"] = edge_filter
def __init__(self, modelparams): """ Initialize model. Args: modelparams (dict): dictionary of parameters for the model Returns: None """ WeightedConformers.__init__(self, modelparams) # get rid of the atom embedding, as we'll be using graph-based # atom features instead of atomic number embeddings delattr(self, "atom_embed") n_convolutions = modelparams["n_convolutions"] dropout_rate = modelparams["dropout_rate"] n_bond_hidden = modelparams["n_bond_hidden"] n_bond_features = modelparams["n_bond_features"] n_atom_basis = modelparams["n_atom_basis"] n_filters = modelparams["n_filters"] trainable_gauss = modelparams["trainable_gauss"] n_gaussians = modelparams["n_gaussians"] cutoff = modelparams["cutoff"] activation = modelparams["activation"] n_atom_hidden = modelparams["n_atom_hidden"] self.convolutions = nn.ModuleList([ MixedSchNetConv(n_atom_hidden=n_atom_hidden, n_filters=n_filters, dropout_rate=dropout_rate, n_bond_hidden=n_bond_hidden, activation=activation) for _ in range(n_convolutions) ]) # for converting distances to features before concatenating with # bond features self.distance_filter = SchNetEdgeFilter( cutoff=cutoff, n_gaussians=n_gaussians, trainable_gauss=trainable_gauss, n_filters=n_filters, dropout_rate=dropout_rate, activation=activation) # for converting bond features to hidden feature vectors self.bond_filter = Sequential( Dense(in_features=n_bond_features, out_features=n_bond_hidden, dropout_rate=dropout_rate), layer_types[activation](), Dense(in_features=n_bond_hidden, out_features=n_bond_hidden, dropout_rate=dropout_rate)) self.atom_filter = Sequential( Dense(in_features=n_atom_basis, out_features=n_atom_hidden, dropout_rate=dropout_rate), layer_types[activation](), Dense(in_features=n_atom_hidden, out_features=n_atom_hidden, dropout_rate=dropout_rate))