def __init__(self, modelparams):
super().__init__()
n_atom_basis = modelparams['n_atom_basis']
n_filters = modelparams['n_filters']
n_gaussians = modelparams['n_gaussians']
trainable_gauss = modelparams.get('trainable_gauss', False)
mol_n_convolutions = modelparams['mol_n_convolutions']
mol_cutoff = modelparams['mol_cutoff']
sys_n_convolutions = modelparams['sys_n_convolutions']
sys_cutoff = modelparams['sys_cutoff']
self.temp_type = modelparams['temp_type']
self.power = modelparams["V_ex_power"]
self.sigma = torch.nn.Parameter(
torch.Tensor([modelparams["V_ex_sigma"]]))
# default predict var
readoutdict = modelparams.get('readoutdict',
get_default_readout(n_atom_basis))
post_readout = modelparams.get('post_readout', None)
dropout_rate = modelparams.get('dropout_rate', 0)
self.atom_embed = nn.Embedding(100, n_atom_basis, padding_idx=0)
self.molecule_convolutions = nn.ModuleList([
SchNetConv(n_atom_basis=n_atom_basis,
n_filters=n_filters,
n_gaussians=n_gaussians,
cutoff=mol_cutoff,
trainable_gauss=trainable_gauss,
dropout_rate=dropout_rate)
for _ in range(mol_n_convolutions)
])
self.system_convolutions = nn.ModuleList([
SchNetConv(n_atom_basis=n_atom_basis,
n_filters=n_filters,
n_gaussians=n_gaussians,
cutoff=sys_cutoff,
trainable_gauss=trainable_gauss,
dropout_rate=dropout_rate)
for _ in range(sys_n_convolutions)
])
# ReadOut
self.atomwisereadout = NodeMultiTaskReadOut(multitaskdict=readoutdict,
post_readout=post_readout)
self.thermo_embed = torch.nn.Linear(1, n_atom_basis)
self.device = None
def __init__(self, modelparams):
"""Constructs a SchNet model.
Args:
modelparams (TYPE): Description
"""
super().__init__()
n_atom_basis = modelparams['n_atom_basis']
n_filters = modelparams['n_filters']
n_gaussians = modelparams['n_gaussians']
n_convolutions = modelparams['n_convolutions']
cutoff = modelparams['cutoff']
trainable_gauss = modelparams.get('trainable_gauss', False)
# default predict var
readoutdict = modelparams.get('readoutdict',
get_default_readout(n_atom_basis))
post_readout = modelparams.get('post_readout', None)
self.atom_embed = nn.Embedding(100, n_atom_basis, padding_idx=0)
self.convolutions = nn.ModuleList([
SchNetConv(n_atom_basis=n_atom_basis,
n_filters=n_filters,
n_gaussians=n_gaussians,
cutoff=cutoff,
trainable_gauss=trainable_gauss)
for _ in range(n_convolutions)
])
# ReadOut
self.atomwisereadout = NodeMultiTaskReadOut(multitaskdict=readoutdict,
post_readout=post_readout)
self.device = None
def __init__(self, modelparams):
"""Constructs a SchNet model.
Args:
modelparams (TYPE): Description
Example:
n_atom_basis = 256
readoutdict = {
"energy_0": [{'name': 'linear', 'param' : { 'in_features': n_atom_basis,
'out_features': int(n_atom_basis / 2)}},
{'name': 'shifted_softplus', 'param': {}},
{'name': 'linear', 'param' : { 'in_features': int(n_atom_basis / 2),
'out_features': 1}}],
"energy_1": [{'name': 'linear', 'param' : { 'in_features': n_atom_basis,
'out_features': int(n_atom_basis / 2)}},
{'name': 'shifted_softplus', 'param': {}},
{'name': 'linear', 'param' : { 'in_features': int(n_atom_basis / 2),
'out_features': 1}}]
}
modelparams = {
'n_atom_basis': n_atom_basis,
'n_filters': 256,
'n_gaussians': 32,
'n_convolutions': 4,
'cutoff': 5.0,
'trainable_gauss': True,
'readoutdict': readoutdict,
'dropout_rate': 0.2
}
model = SchNet(modelparams)
"""
nn.Module.__init__(self)
n_atom_basis = modelparams["n_atom_basis"]
n_filters = modelparams["n_filters"]
n_gaussians = modelparams["n_gaussians"]
n_convolutions = modelparams["n_convolutions"]
cutoff = modelparams["cutoff"]
trainable_gauss = modelparams.get("trainable_gauss", False)
dropout_rate = modelparams.get("dropout_rate", DEFAULT_DROPOUT_RATE)
self.atom_embed = nn.Embedding(100, n_atom_basis, padding_idx=0)
readoutdict = modelparams.get("readoutdict",
get_default_readout(n_atom_basis))
post_readout = modelparams.get("post_readout", None)
# convolutions
self.convolutions = nn.ModuleList([
SchNetConv(
n_atom_basis=n_atom_basis,
n_filters=n_filters,
n_gaussians=n_gaussians,
cutoff=cutoff,
trainable_gauss=trainable_gauss,
dropout_rate=dropout_rate,
) for _ in range(n_convolutions)
])
# ReadOut
self.atomwisereadout = NodeMultiTaskReadOut(multitaskdict=readoutdict,
post_readout=post_readout)
self.device = None
def __init__(self, modelparams):
"""Constructs a SchNet-Like model using a conformer representation.
Args:
modelparams (dict): dictionary of parameters for model. All
are the same as in SchNet, except for `mol_fp_layers`,
which describes how to convert atomic fingerprints into
a single molecular fingerprint.
Example:
n_atom_basis = 256
mol_basis = 512
# all the atomic fingerprints get added together, then go through the network created
# by `mol_fp_layers` to turn into a molecular fingerprint
mol_fp_layers = [{'name': 'linear', 'param' : { 'in_features': n_atom_basis,
'out_features': int((n_atom_basis + mol_basis)/2)}},
{'name': 'shifted_softplus', 'param': {}},
{'name': 'linear', 'param' : { 'in_features': int((n_atom_basis + mol_basis)/2),
'out_features': mol_basis}}]
readoutdict = {
"covid": [{'name': 'linear', 'param' : { 'in_features': mol_basis,
'out_features': int(mol_basis / 2)}},
{'name': 'shifted_softplus', 'param': {}},
{'name': 'linear', 'param' : { 'in_features': int(mol_basis / 2),
'out_features': 1}},
{'name': 'sigmoid', 'param': {}}],
}
# dictionary to tell you what to do with the Boltzmann factors
# ex. 1:
boltzmann_dict = {"type": "multiply"}
# ex. 2
boltzmann_layers = [{'name': 'linear', 'param': {'in_features': mol_basis + 1,
'out_features': mol_basis}},
{'name': 'shifted_softplus', 'param': {}},
{'name': 'linear', 'param': {'in_features': mol_basis,
'out_features': mol_basis}}]
boltzmann_dict = {"type": "layers", "layers": boltzmann_layers}
modelparams = {
'n_atom_basis': n_atom_basis,
'n_filters': 256,
'n_gaussians': 32,
'n_convolutions': 4,
'cutoff': 5.0,
'trainable_gauss': True,
'readoutdict': readoutdict,
'mol_fp_layers': mol_fp_layers,
'boltzmann_dict': boltzmann_dict
'dropout_rate': 0.2
}
model = WeightedConformers(modelparams)
"""
nn.Module.__init__(self)
n_atom_basis = modelparams["n_atom_basis"]
n_filters = modelparams["n_filters"]
n_gaussians = modelparams["n_gaussians"]
n_convolutions = modelparams["n_convolutions"]
cutoff = modelparams["cutoff"]
trainable_gauss = modelparams.get("trainable_gauss", False)
dropout_rate = modelparams.get("dropout_rate", DEFAULT_DROPOUT_RATE)
self.atom_embed = nn.Embedding(100, n_atom_basis, padding_idx=0)
# convolutions
self.convolutions = nn.ModuleList(
[
SchNetConv(
n_atom_basis=n_atom_basis,
n_filters=n_filters,
n_gaussians=n_gaussians,
cutoff=cutoff,
trainable_gauss=trainable_gauss,
dropout_rate=dropout_rate,
)
for _ in range(n_convolutions)
]
)
# extra features to consider
self.extra_feats = modelparams.get("extra_features")
self.ext_feat_types = modelparams.get("ext_feat_types")
mol_fp_layers = modelparams["mol_fp_layers"]
readoutdict = modelparams["readoutdict"]
boltzmann_dict = modelparams["boltzmann_dict"]
# the nn that converts atomic finerprints to a molecular fp
self.mol_fp_nn = construct_sequential(mol_fp_layers)
# create a module that lets a molecular fp interact with the
# conformer's boltzmann weight to give a final molecular fp
self.boltz_nns = self.make_boltz_nn(boltzmann_dict)
self.head_pool = boltzmann_dict.get("head_pool", "concatenate")
# the readout acts on this final molceular fp
self.readout = NodeMultiTaskReadOut(multitaskdict=readoutdict)
# whether this is a classifier
self.classifier = modelparams["classifier"]
# whether to embed fingerprints or just use external features
self.use_mpnn = modelparams.get("use_mpnn", True)