def __init__(self, input_layers):
"""
Args:
input_layers (list[dict]): instructions for
making the input layers applied to the node
and edge features.
Returns:
None
"""
nn.Module.__init__(self)
# remove bias from linear layers if there
new_layers = remove_bias(input_layers)
self.input = construct_sequential(new_layers)
def __init__(self, output_layers):
"""
Args:
output_layers (list[dict]): instructions for
making the output layers applied after the
initial node features get concatenated with
the edge-turned-node updated features.
Returns:
None
"""
nn.Module.__init__(self)
# remove bias from linear layers if there
new_layers = remove_bias(output_layers)
self.output = construct_sequential(new_layers)
def __init__(self, update_layers):
super(AuTopologyConv, self).__init__()
"""
Args:
update_layers (dict): dictionary of layers to apply after the convolution
Returns:
None
Example:
update_layers = [{'name': 'linear', 'param' : {'in_features': 256,
'out_features': 256}},
{'name': 'tanh', 'param': {}},
{'name': 'linear', 'param' : {'in_features': 256,
'out_features': 256}},
{'name': 'tanh', 'param': {}}
"""
self.update_function = construct_sequential(update_layers)
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)
def make_boltz_nn(self, boltzmann_dict):
"""
Make the section of the network that creates weights for each
conformer, which may or may not be equal to the statistical
boltzmann weights.
Args:
boltzmann_dict (dict): dictionary with information about
this section of the network.
Returns:
networks (nn.ModuleList): list of networks that get applied
to the conformer fingerprints to aggregate them. If
it contains more than one network, the different fingerprints
produced will either be averaged or concatenated at the end.
"""
networks = nn.ModuleList([])
# if you just want to multiply the boltmzann weight by each conformer
# fingerprint, return nothing
if boltzmann_dict["type"] == "multiply":
return [None]
# if you supply a dictionary of type `layers`, then the dictionary
# under the key `layers` will be used to create the corresponding
# network
elif boltzmann_dict["type"] == "layers":
layers = boltzmann_dict["layers"]
networks.append(construct_sequential(layers))
# if you ask for some sort of attention network, then make one such
# network for each of the number of heads
elif "attention" in boltzmann_dict["type"]:
if boltzmann_dict["type"] == "attention":
module = ConfAttention
elif boltzmann_dict["type"] == "linear_attention":
module = LinearConfAttention
else:
raise NotImplementedError
# how many attention heads
num_heads = boltzmann_dict.get("num_heads", 1)
# whether to just use equal weights and not learnable weights
# (useful for ablation studies)
equal_weights = boltzmann_dict.get("equal_weights", False)
# what function to use to convert the alpha_ij to probabilities
prob_func = boltzmann_dict.get("prob_func", 'softmax')
# add a network for each head
for _ in range(num_heads):
mol_basis = boltzmann_dict["mol_basis"]
boltz_basis = boltzmann_dict["boltz_basis"]
final_act = boltzmann_dict["final_act"]
networks.append(module(mol_basis=mol_basis,
boltz_basis=boltz_basis,
final_act=final_act,
equal_weights=equal_weights,
prob_func=prob_func))
return networks