Models

Modeling

GIN

class molgraph.models.GIN(tensorflow.keras.layers.Layer)[source]

Graph isomorphism network (GIN), with or without edge features.

Implementation based on Xu et al. (2019) 1 or Hu et al. (2020) 2.

Example usage:

>>> # Obtain GraphTensor
>>> graph_tensor = molgraph.GraphTensor(
...     sizes=[2, 3],
...     node_feature=[[1., 0.], [1., 0.], [1., 0.], [1., 0.], [0., 1.]],
...     edge_feature=[[1., 0.], [0., 1.], [0., 1.], [0., 1.],
...                   [1., 0.], [0., 1.], [1., 0.], [0., 1.]],
...     edge_src=[1, 0, 3, 4, 2, 4, 3, 2],
...     edge_dst=[0, 1, 2, 2, 3, 3, 4, 4],
... )
>>> # Build Functional model
>>> inputs = tf.keras.layers.Input(type_spec=graph_tensor.spec)
>>> # (n_nodes, n_features) -> (n_nodes, n_steps, n_features)
>>> x = molgraph.models.GIN(steps=4, units=32, merge_mode='stack')(inputs)
>>> # (n_nodes, n_steps, n_features) -> (n_graphs, n_steps, n_features)
>>> outputs = molgraph.layers.Readout()(x)
>>> gin_model = tf.keras.Model(inputs, outputs)
>>> graph_embeddings = gin_model.predict(graph_tensor, verbose=0)
>>> graph_embeddings.shape
(2, 5, 32)

Parameters

steps (int) – Number of message passing steps. Default to 4.
units (int, None) – Number of hidden units of the message passing. Default to None.
merge_mode (str) – How the node embeddings should be merged. Either of ‘concat’, ‘stack’, ‘sum’, ‘mean’, ‘max’ or ‘min’. Default to ‘stack’.
use_edge_features (bool) – Whether or not to use edge features. Default to False.
apply_relu_activation (bool) – Whether to apply relu activation before aggregation step. Only relevant if use_edge_features is set to True. Default to False.
normalization – (None, str, bool): Whether to apply layer normalization to the output. If batch normalization is desired, pass ‘batch_norm’. Default to True.
residual – (bool) Whether to add skip connection to the output. Default to True.
dropout – (float, None): Dropout applied to the output of the layer. Default to None.
activation (tf.keras.activations.Activation, callable, str, None) – Activation function applied to the projections. Default to ‘relu’.
use_bias (bool) – Whether the layer should use biases. Default to True.
kernel_initializer (tf.keras.initializers.Initializer, str) – Initializer function for the kernels. If None, tf.keras.initializers.TruncatedNormal(stddev=0.005) will be used. Default to None.
bias_initializer (tf.keras.initializers.Initializer, str) – Initializer function for the biases. Default to tf.keras.initializers.Constant(0.).
kernel_regularizer (tf.keras.regularizers.Regularizer, None) – Regularizer function applied to the kernels. Default to None.
bias_regularizer (tf.keras.regularizers.Regularizer, None) – Regularizer function applied to the biases. Default to None.
activity_regularizer (tf.keras.regularizers.Regularizer, None) – Regularizer function applied to the final output of the layer. Default to None.
kernel_constraint (tf.keras.constraints.Constraint, None) – Constraint function applied to the kernels. Default to None.
bias_constraint (tf.keras.constraints.Constraint, None) – Constraint function applied to the biases. Default to None.

References

1: https://arxiv.org/pdf/1810.00826.pdf
2: https://arxiv.org/pdf/1905.12265.pdf

build(input_shape)[source]

Creates the variables of the layer (for subclass implementers).

This is a method that implementers of subclasses of Layer or Model can override if they need a state-creation step in-between layer instantiation and layer call. It is invoked automatically before the first execution of call().

This is typically used to create the weights of Layer subclasses (at the discretion of the subclass implementer).

Parameters: input_shape – Instance of TensorShape, or list of instances of TensorShape if the layer expects a list of inputs (one instance per input).

call(tensor)[source]

Defines the computation from inputs to outputs.

This method should not be called directly, but indirectly via __call__(). Upon first call, the layer is automatically built via build().

Parameters: tensor (GraphTensor) – Input to the layer.
Returns: A GraphTensor with updated node features.
Return type: GraphTensor

get_config()[source]

Returns the config of the layer.

A layer config is a Python dictionary (serializable) containing the configuration of a layer. The same layer can be reinstantiated later (without its trained weights) from this configuration.

The config of a layer does not include connectivity information, nor the layer class name. These are handled by Network (one layer of abstraction above).

Note that get_config() does not guarantee to return a fresh copy of dict every time it is called. The callers should make a copy of the returned dict if they want to modify it.

Returns: Python dictionary.

classmethod from_config(config)[source]

Creates a layer from its config.

This method is the reverse of get_config, capable of instantiating the same layer from the config dictionary. It does not handle layer connectivity (handled by Network), nor weights (handled by set_weights).

Parameters: config – A Python dictionary, typically the output of get_config.
Returns: A layer instance.

MPNN

class molgraph.models.MPNN(tensorflow.keras.layers.Layer)[source]

Message passing neural network (MPNN) with weight tying.

Implementation is based on Gilmer et al. (2017) 3.

Example usage:

>>> # Obtain GraphTensor
>>> graph_tensor = molgraph.GraphTensor(
...     sizes=[2, 3],
...     node_feature=[[1., 0.], [1., 0.], [1., 0.], [1., 0.], [0., 1.]],
...     edge_feature=[[1., 0.], [0., 1.], [0., 1.], [0., 1.],
...                   [1., 0.], [0., 1.], [1., 0.], [0., 1.]],
...     edge_src=[1, 0, 3, 4, 2, 4, 3, 2],
...     edge_dst=[0, 1, 2, 2, 3, 3, 4, 4],
... )
>>> # Build Functional model
>>> inputs = tf.keras.layers.Input(type_spec=graph_tensor.spec)
>>> x = molgraph.models.MPNN(units=32, steps=4, name='mpnn')(inputs)
>>> x = molgraph.layers.SetGatherReadout(name='readout')(x)
>>> outputs = tf.keras.layers.Dense(10, activation='sigmoid')(x)
>>> mpnn_classifier = tf.keras.Model(inputs, outputs)
>>> # Make predictions
>>> preds = mpnn_classifier.predict(graph_tensor, verbose=0)
>>> preds.shape
(2, 10)

Parameters

units (int, None) – Number of hiden units of the message passing. These include the dense layers associated with the message functions, and GRU cells associated with the update functions. If None, hidden units are set to the input dimension. Default to None.
steps (int) – Number of message passing steps. Default to 4.
residual – (bool) Whether to add skip connection to the output of each step. Default to True.
dropout – (float, None): Dropout applied to the output of step. Default to None.
message_kwargs (dict, None) – An optional dictionary of parameters which can be passed to the dense layers of the message functions. Note: as units is already specified, it will be dropped from the dict (if it exists there). If None, an empty dict will be passed. Default to None.
update_kwargs (dict, None) – An optional dictionary of parameters which can be passed to the GRUCells of the update functions. Note: as units is already specified, it will be dropped from the dict (if it exists there). If None, an empty dict will be passed. Default to None.

References

3: https://arxiv.org/pdf/1704.01212.pdf

build(input_shape)[source]

Creates the variables of the layer (for subclass implementers).

This is a method that implementers of subclasses of Layer or Model can override if they need a state-creation step in-between layer instantiation and layer call. It is invoked automatically before the first execution of call().

This is typically used to create the weights of Layer subclasses (at the discretion of the subclass implementer).

Parameters: input_shape – Instance of TensorShape, or list of instances of TensorShape if the layer expects a list of inputs (one instance per input).

call(tensor)[source]

Defines the computation from inputs to outputs.

This method should not be called directly, but indirectly via __call__(). Upon first call, the layer is automatically built via build().

Parameters: tensor (GraphTensor) – Input to the layer.
Returns: A GraphTensor with updated node features.
Return type: GraphTensor

get_config()[source]

Returns the config of the layer.

A layer config is a Python dictionary (serializable) containing the configuration of a layer. The same layer can be reinstantiated later (without its trained weights) from this configuration.

The config of a layer does not include connectivity information, nor the layer class name. These are handled by Network (one layer of abstraction above).

Note that get_config() does not guarantee to return a fresh copy of dict every time it is called. The callers should make a copy of the returned dict if they want to modify it.

Returns: Python dictionary.

classmethod from_config(config)[source]

Creates a layer from its config.

This method is the reverse of get_config, capable of instantiating the same layer from the config dictionary. It does not handle layer connectivity (handled by Network), nor weights (handled by set_weights).

Parameters: config – A Python dictionary, typically the output of get_config.
Returns: A layer instance.

DMPNN

class molgraph.models.DMPNN(tensorflow.keras.layers.Layer)[source]

Directed message passing neural network (DMPNN).

Implementation based on Yang et al. (2019) 4.

As of now, DMPNN only works on (sub)graphs with at least one edge/bond. If your dataset consists of molecules with single atoms, please add self loops: molgraph.chemistry.MolecularGraphEncoder(..., self_loops=True).

Example usage:

>>> # Obtain GraphTensor
>>> graph_tensor = molgraph.GraphTensor(
...     sizes=[2, 3],
...     node_feature=[[1., 0.], [1., 0.], [1., 0.], [1., 0.], [0., 1.]],
...     edge_feature=[[1., 0.], [0., 1.], [0., 1.], [0., 1.],
...                   [1., 0.], [0., 1.], [1., 0.], [0., 1.]],
...     edge_src=[1, 0, 3, 4, 2, 4, 3, 2],
...     edge_dst=[0, 1, 2, 2, 3, 3, 4, 4],
... )
>>> # Build Functional model
>>> inputs = tf.keras.layers.Input(type_spec=graph_tensor.spec)
>>> x = molgraph.models.DMPNN(units=32, name='dmpnn')(inputs)
>>> x = molgraph.layers.Readout(name='readout')(x)
>>> outputs = tf.keras.layers.Dense(10, activation='sigmoid')(x)
>>> dmpnn_classifier = tf.keras.Model(inputs, outputs)
>>> # Make predictions
>>> preds = dmpnn_classifier.predict(graph_tensor, verbose=0)
>>> preds.shape
(2, 10)

Parameters

steps (int) – Number of message passing steps. Default to 4.
units (int, None) – Number of hidden units of the message passing. Default to None.
normalization – (None, str, bool): Whether to apply layer normalization to the output. If batch normalization is desired, pass ‘batch_norm’. Default to True.
residual – (bool) Whether to add skip connection to the output. Default to True.
dropout – (float, None): Dropout applied to the output of the layer. Default to None.
activation (tf.keras.activations.Activation, callable, str, None) – Activation function applied to the projections. Default to ‘relu’.
use_bias (bool) – Whether the layer should use biases. Default to True.
kernel_initializer (tf.keras.initializers.Initializer, str) – Initializer function for the kernels. If None, tf.keras.initializers.TruncatedNormal(stddev=0.005) will be used. Default to None.
bias_initializer (tf.keras.initializers.Initializer, str) – Initializer function for the biases. Default to tf.keras.initializers.Constant(0.).
kernel_regularizer (tf.keras.regularizers.Regularizer, None) – Regularizer function applied to the kernels. Default to None.
bias_regularizer (tf.keras.regularizers.Regularizer, None) – Regularizer function applied to the biases. Default to None.
activity_regularizer (tf.keras.regularizers.Regularizer, None) – Regularizer function applied to the final output of the layer. Default to None.
kernel_constraint (tf.keras.constraints.Constraint, None) – Constraint function applied to the kernels. Default to None.
bias_constraint (tf.keras.constraints.Constraint, None) – Constraint function applied to the biases. Default to None.

References

4: https://arxiv.org/abs/1904.01561

build(input_shape)[source]

Creates the variables of the layer (for subclass implementers).

This is a method that implementers of subclasses of Layer or Model can override if they need a state-creation step in-between layer instantiation and layer call. It is invoked automatically before the first execution of call().

This is typically used to create the weights of Layer subclasses (at the discretion of the subclass implementer).

Parameters: input_shape – Instance of TensorShape, or list of instances of TensorShape if the layer expects a list of inputs (one instance per input).

call(tensor)[source]

Defines the computation from inputs to outputs.

This method should not be called directly, but indirectly via __call__(). Upon first call, the layer is automatically built via build().

Parameters: tensor (GraphTensor) – Input to the layer.
Returns: A GraphTensor with updated node features.
Return type: GraphTensor

get_config()[source]

Returns the config of the layer.

A layer config is a Python dictionary (serializable) containing the configuration of a layer. The same layer can be reinstantiated later (without its trained weights) from this configuration.

The config of a layer does not include connectivity information, nor the layer class name. These are handled by Network (one layer of abstraction above).

Note that get_config() does not guarantee to return a fresh copy of dict every time it is called. The callers should make a copy of the returned dict if they want to modify it.

Returns: Python dictionary.

classmethod from_config(config)[source]

Creates a layer from its config.

This method is the reverse of get_config, capable of instantiating the same layer from the config dictionary. It does not handle layer connectivity (handled by Network), nor weights (handled by set_weights).

Parameters: config – A Python dictionary, typically the output of get_config.
Returns: A layer instance.

Interpretability

GradientActivationMapping

class molgraph.models.GradientActivationMapping(tensorflow.Module)[source]

Gradient activation mapping.

Implementation is based on Pope et al. 5.

Alias: GradientActivation

Example usage:

>>> encoder = molgraph.chemistry.MolecularGraphEncoder(
...     atom_encoder=molgraph.chemistry.Featurizer([
...         molgraph.chemistry.features.Symbol(),
...         molgraph.chemistry.features.Hybridization()
...     ])
... )
>>> esol = molgraph.chemistry.datasets.get('esol')
>>> esol['train']['x'] = encoder(esol['train']['x'])
>>> esol['test']['x'] = encoder(esol['test']['x'])
...
>>> gnn_model = tf.keras.Sequential([
...     molgraph.layers.GCNConv(units=128, name='gcn_conv_1'),
...     molgraph.layers.GCNConv(units=128, name='gcn_conv_2'),
...     molgraph.layers.GCNConv(units=128, name='gcn_conv_3'),
...     molgraph.layers.Readout('mean'),
...     tf.keras.layers.Dense(units=512),
...     tf.keras.layers.Dense(units=1)
... ])
...
>>> gnn_model.compile(optimizer='adam', loss='mse')
>>> _ = gnn_model.fit(
...     esol['train']['x'], esol['train']['y'], epochs=10, verbose=0)
...
>>> gam_model = molgraph.models.GradientActivationMapping(
...     model=gnn_model,
...     discard_negative_values=False,
... )
>>> maps = gam_model(esol['test']['x'].separate())

References

5: Pope et al. https://ieeexplore.ieee.org/document/8954227

__call__(x, y=None)[source]: Call self as a function.

SaliencyMapping

class molgraph.models.SaliencyMapping(tensorflow.Module)[source]

Vanilla saliency mapping.

Alias: Saliency

Example usage:

>>> encoder = molgraph.chemistry.MolecularGraphEncoder(
...     atom_encoder=molgraph.chemistry.Featurizer([
...         molgraph.chemistry.features.Symbol(),
...         molgraph.chemistry.features.Hybridization()
...     ])
... )
...
>>> esol = molgraph.chemistry.datasets.get('esol')
>>> esol['train']['x'] = encoder(esol['train']['x'])
>>> esol['test']['x'] = encoder(esol['test']['x'])
...
>>> gnn_model = tf.keras.Sequential([
...     molgraph.layers.GCNConv(units=128, name='gcn_conv_1'),
...     molgraph.layers.GCNConv(units=128, name='gcn_conv_2'),
...     molgraph.layers.GCNConv(units=128, name='gcn_conv_3'),
...     molgraph.layers.Readout('mean'),
...     tf.keras.layers.Dense(units=512),
...     tf.keras.layers.Dense(units=1)
... ])
...
>>> gnn_model.compile(optimizer='adam', loss='mse')
>>> _ = gnn_model.fit(
...     esol['train']['x'], esol['train']['y'], epochs=10, verbose=0)
...
>>> saliency = molgraph.models.SaliencyMapping(model=gnn_model)
>>> saliency_maps = saliency(esol['test']['x'].separate())

__call__(x, y=None)[source]: Call self as a function.

IntegratedSaliencyMapping

class molgraph.models.IntegratedSaliencyMapping(tensorflow.Module)[source]

Integrated saliency mapping.

Alias: IntegratedSaliency

Example usage:

>>> encoder = molgraph.chemistry.MolecularGraphEncoder(
...     atom_encoder=molgraph.chemistry.Featurizer([
...         molgraph.chemistry.features.Symbol(),
...         molgraph.chemistry.features.Hybridization()
...     ])
... )
>>> bbbp = molgraph.chemistry.datasets.get('bbbp')
>>> bbbp['train']['x'] = encoder(bbbp['train']['x'])
>>> bbbp['test']['x'] = encoder(bbbp['test']['x'])
...
>>> gnn_model = tf.keras.Sequential([
...     molgraph.layers.GCNConv(units=128, name='gcn_conv_1'),
...     molgraph.layers.GCNConv(units=128, name='gcn_conv_2'),
...     molgraph.layers.GCNConv(units=128, name='gcn_conv_3'),
...     molgraph.layers.Readout('mean'),
...     tf.keras.layers.Dense(units=512),
...     tf.keras.layers.Dense(units=1, activation='sigmoid')
... ])
...
>>> gnn_model.compile(optimizer='adam', loss='mse')
>>> _ = gnn_model.fit(
...     bbbp['train']['x'], bbbp['train']['y'], epochs=10, verbose=0)
...
>>> saliency = molgraph.models.IntegratedSaliencyMapping(model=gnn_model)
>>> saliency_maps = saliency(bbbp['test']['x'].separate())

__call__(x, y=None)[source]: Call self as a function.

SmoothGradSaliencyMapping

class molgraph.models.SmoothGradSaliencyMapping(tensorflow.Module)[source]

Smooth-gradient saliency mapping.

Alias: SmoothGradSaliency

Example usage:

>>> encoder = molgraph.chemistry.MolecularGraphEncoder(
...     atom_encoder=molgraph.chemistry.Featurizer([
...         molgraph.chemistry.features.Symbol(),
...         molgraph.chemistry.features.Hybridization()
...     ])
... )
>>> esol = molgraph.chemistry.datasets.get('esol')
>>> esol['train']['x'] = encoder(esol['train']['x'])
>>> esol['test']['x'] = encoder(esol['test']['x'])
...
>>> gnn_model = tf.keras.Sequential([
...     molgraph.layers.GCNConv(units=128, name='gcn_conv_1'),
...     molgraph.layers.GCNConv(units=128, name='gcn_conv_2'),
...     molgraph.layers.GCNConv(units=128, name='gcn_conv_3'),
...     molgraph.layers.Readout('mean'),
...     tf.keras.layers.Dense(units=512),
...     tf.keras.layers.Dense(units=1)
... ])
...
>>> gnn_model.compile(optimizer='adam', loss='mse')
>>> _ = gnn_model.fit(
...     esol['train']['x'], esol['train']['y'], epochs=10, verbose=0)
...
>>> saliency = molgraph.models.SmoothGradSaliencyMapping(model=gnn_model)
>>> saliency_maps = saliency(esol['test']['x'].separate())

__call__(x, y=None)[source]: Call self as a function.

Pretraining

GraphMasking

class molgraph.models.GraphMasking(tensorflow.keras.Model)[source]

Masked Graph Modeling (MGM) inspired by Masked Language Modeling (MLM).

See e.g. Hu et al. 6 and Devlin et al. 7.

The encoder part of the MGM model is the model to be finetuned for downstream modeling. In this MGM model, masking layer(s) will be prepended to randomly mask out node and/or edge features of a certain rate. Furthermore, a classification layer is appended to the encoder to produce multi-class predictions; with the objective to predict the masked node and/or edge features.

Currently, this MGM model only support Tokenized molecular graphs which are embedded via EmbeddingLookup. A Featurized molecular graph (without EmbeddingLookup) is less straight forward to implement, though it could be implemented by e.g. creating random vectors for each masked node/edge.

Example usage:

>>> # Replace this graph_tensor with a large dataset of graphs
>>> graph_tensor = molgraph.GraphTensor(
...     sizes=[5],
...     edge_src=[1, 4, 0, 2, 3, 1, 1, 0],
...     edge_dst=[0, 0, 1, 1, 1, 2, 3, 4],
...     node_feature=['Sym:C|Hyb:SP3', 'Sym:C|Hyb:SP2', 'Sym:O|Hyb:SP2',
...                   'Sym:O|Hyb:SP2', 'Sym:N|Hyb:SP3'],
...     edge_feature=['BonTyp:SINGLE|Rot:1', 'BonTyp:SINGLE|Rot:0',
...                   'BonTyp:SINGLE|Rot:1', 'BonTyp:DOUBLE|Rot:0',
...                   'BonTyp:SINGLE|Rot:0', 'BonTyp:DOUBLE|Rot:0',
...                   'BonTyp:SINGLE|Rot:0', 'BonTyp:SINGLE|Rot:0'],
... )
>>> node_embedding = molgraph.layers.NodeEmbeddingLookup(
...    32, mask_token='[MASK]')
>>> edge_embedding = molgraph.layers.EdgeEmbeddingLookup(
...    32, mask_token='[MASK]')
>>> node_embedding.adapt(graph_tensor)
>>> edge_embedding.adapt(graph_tensor)
>>> graph_transformer_encoder = tf.keras.Sequential([
...     node_embedding,
...     edge_embedding,
...     molgraph.layers.GTConv(units=32),
...     molgraph.layers.GTConv(units=32),
...     molgraph.layers.GTConv(units=32),
... ])
>>> pretraining_model = molgraph.models.GraphMasking(
...     graph_transformer_encoder,
...     node_feature_masking_rate=0.15,
...     edge_feature_masking_rate=0.15
... )
>>> pretraining_model.compile(
...     optimizer=tf.keras.optimizers.Adam(1e-4),
...     metrics=[
...         tf.keras.metrics.SparseCategoricalAccuracy(name='sparse_acc'),
...     ]
... )
>>> _ = pretraining_model.fit(graph_tensor, epochs=1, verbose=0)
>>> metric_values = pretraining_model.evaluate(graph_tensor, verbose=0)
>>> graph_transformer_encoder.save( 
...     '/tmp/my_pretrained_encoder_model'
... )
>>> loaded_encoder = tf.saved_model.load( 
...     '/tmp/my_pretrained_encoder_model'
... )

Parameters

encoder (tf.keras.Model) – The model to be pretrained and used for downstream tasks. The first layer(s) of the model should embed (via EmbeddingLookup) (later on, possibly masked) node/edge embeddings from tokenized molecular graphs (obtain via Tokenizer).
node_feature_decoder (None, tf.keras.layers.Layer) – Optionally supply a decoder which turns the node embeddings into a softmaxed prediction. If None, a linear transdformation followed by a softmax activation will be used. Default to None.
edge_feature_decoder (None, tf.keras.layers.Layer) – Optionally supply a decoder which turns the edge embeddings into a softmaxed prediction. If None, a linear transdformation followed by a softmax activation will be used. Default to None.
node_feature_masking_rate (None, float) – The rate at which node features should be masked. If None, or 0.0, no masking will be performed on the node features. Default to 0.15.
edge_feature_masking_rate (None, float) – The rate at which edge features should be masked. If None, or 0.0, no masking will be performed on the edge features. Default to 0.15.

References

6: https://arxiv.org/pdf/1905.12265.pdf
7: https://arxiv.org/pdf/1810.04805.pdf

compile(optimizer, loss=None, metrics=None, node_feature_loss_weight=None, edge_feature_loss_weight=None, *args, **kwargs)[source]

Configures the model for training.

Parameters

optimizer (tf.keras.optimizers.Optimizer) – The optimizer to use for training.
loss (None, tf.keras.losses.Loss) – The loss function to use. If None, tf.keras.losses.SparseCategoricalCrossentropy will be used. If a custom loss function is used, make sure it deals with sparse labels (i.e. integer encoding and not one-hot encoding). Default to None.
metrics (None, list[tf.keras.metrics.Metric]) – The metrics to use. Default to None.
node_feature_loss_weight (None, float) – The weight to be applied to the node feature prediction loss. If None, the node feature loss will be multiplied by 1. Default to None.
edge_feature_loss_weight (None, float) – THe wieght to be applied to the edge feature prediction loss. If None, the edge feature loss will be multiplied by 1. Default to None.
*args – See tf.keras.Model.compile documentation.
**kwargs – See tf.keras.Model.compile documentation.

fit(x, y=None, batch_size=None, epochs=1, *args, **kwargs)[source]

Trains the autoencoder for a fixed number of epochs.

Parameters

x (GraphTensor, tf.data.Dataset) – The input data. Either a GraphTensor instance or a tf.data.Dataset of `GraphTensor`s.
y (None) – Target data, which will be ignored as the target can be obtained from x.
batch_size (int, None) – Number of samples per gradient update. If None, 32 will be used. Default to None.
epochs (int) – Number of iterations over all subgraphs (molecular graphs) of the GraphTensor instance or tf.data.Dataset instance. Default to 1.
*args – See tf.keras.Model.fit documentaton.
**kwargs – See tf.keras.Model.fit documentation.

Returns

A History object containing e.g. training loss values and metric values.

evaluate(x, y=None, batch_size=None, *args, **kwargs)[source]

Evaluates the autoencoder.

Parameters

x (GraphTensor, tf.data.Dataset) – The input data; either a GraphTensor instance or a tf.data.Dataset of `GraphTensor`s.
y (None) – Target data; which will be ignored as the target can be obtained from x.
batch_size (int, None) – Number of samples per batch of computation. If None, 32 will be used. Default to None.
*args – See tf.keras.Model.evaluate documentaton.
**kwargs – See tf.keras.Model.evaluate documentation.

Returns

Average loss values (e.g. reconstruction loss and kl loss).

predict(x, batch_size=None, *args, **kwargs)[source]

Generates outputs of the autoencoder.

Parameters

x (GraphTensor, tf.data.Dataset) – The input data; either a GraphTensor instance or a tf.data.Dataset of `GraphTensor`s.
batch_size (int, None) – Number of samples per batch of computation. If None, 32 will be used. Default to None.
*args – See tf.keras.Model.evaluate documentaton.
**kwargs – See tf.keras.Model.evaluate documentation.

Returns

tf.Tensor or tf.RaggedTensor of edge scores, corresponding to the inputted GraphTensor instance.

classmethod from_config(config, custom_objects=None)[source]

Creates a layer from its config.

This method is the reverse of get_config, capable of instantiating the same layer from the config dictionary. It does not handle layer connectivity (handled by Network), nor weights (handled by set_weights).

Parameters: config – A Python dictionary, typically the output of get_config.
Returns: A layer instance.

get_config()[source]

Returns the config of the Model.

Config is a Python dictionary (serializable) containing the configuration of an object, which in this case is a Model. This allows the Model to be be reinstantiated later (without its trained weights) from this configuration.

Note that get_config() does not guarantee to return a fresh copy of dict every time it is called. The callers should make a copy of the returned dict if they want to modify it.

Developers of subclassed Model are advised to override this method, and continue to update the dict from super(MyModel, self).get_config() to provide the proper configuration of this Model. The default config will return config dict for init parameters if they are basic types. Raises NotImplementedError when in cases where a custom get_config() implementation is required for the subclassed model.

Returns: Python dictionary containing the configuration of this Model.

GraphAutoEncoder

class molgraph.models.GraphAutoEncoder(tensorflow.keras.Model)[source]

Graph AutoEncoder (GAE) based on Kipf and Welling 8.

Parameters

encoder (tf.keras.layers.Layer) – The encoder part of the autoencoder. The encoder could be any of the graph neural neural network layers provided by molgraph; e.g., GCNConv, GINConv or GATv2Conv.
decoder (tf.keras.layers.Layer) – The decoder part of the autoencoder. If None is passed, the decoder used is DotProductIncident. Default to None.
negative_graph_sampler (callable) – A function which samples negative graphs. It takes as input the current GraphTensor instance and produces a new GraphTensor instance with negative edges. If None, NegativeGraphSampler will be used. Default to None.
balanced_class_weighting (bool) – Whether balanced class weighting should be performed. For instance, if there are twice as many negative edges, should each negative edge be weighted 0.5 for the loss?

Example usage:

>>> # Replace this graph_tensor with a large dataset of graphs
>>> # Also accepts featurized graph tensors (via `chemistry.Featurizer`)
>>> graph_tensor = molgraph.GraphTensor(
...     sizes=[5],
...     edge_src=[1, 4, 0, 2, 3, 1, 1, 0],
...     edge_dst=[0, 0, 1, 1, 1, 2, 3, 4],
...     node_feature=['Sym:C|Hyb:SP3', 'Sym:C|Hyb:SP2', 'Sym:O|Hyb:SP2',
...                   'Sym:O|Hyb:SP2', 'Sym:N|Hyb:SP3'],
...     edge_feature=['BonTyp:SINGLE|Rot:1', 'BonTyp:SINGLE|Rot:0',
...                   'BonTyp:SINGLE|Rot:1', 'BonTyp:DOUBLE|Rot:0',
...                   'BonTyp:SINGLE|Rot:0', 'BonTyp:DOUBLE|Rot:0',
...                   'BonTyp:SINGLE|Rot:0', 'BonTyp:SINGLE|Rot:0'],
... )
>>> node_embedding = molgraph.layers.NodeEmbeddingLookup(
...    32, mask_token='[MASK]')
>>> edge_embedding = molgraph.layers.EdgeEmbeddingLookup(
...    32, mask_token='[MASK]')
>>> node_embedding.adapt(graph_tensor)
>>> edge_embedding.adapt(graph_tensor)
>>> # Obtain GraphAutoEncoder model
>>> encoder = tf.keras.Sequential([
...     tf.keras.layers.Input(type_spec=graph_tensor.spec),
...     node_embedding,
...     edge_embedding,
...     molgraph.layers.GATv2Conv(128),
...     molgraph.layers.GATv2Conv(128),
...     molgraph.layers.GATv2Conv(128),
... ])
>>> decoder = molgraph.layers.DotProductIncident(normalize=True)
>>> gae = molgraph.models.GraphAutoEncoder(encoder, decoder)
>>> gae.compile('adam')
>>> _ = gae.fit(graph_tensor, batch_size=32, epochs=50, verbose=0)
>>> reconstruction_loss = gae.evaluate(graph_tensor, verbose=0)
>>> encoder.save( 
...     '/tmp/my_pretrained_encoder_model'
... )
>>> loaded_encoder = tf.saved_model.load( 
...     '/tmp/my_pretrained_encoder_model'
... )

References

8: https://arxiv.org/pdf/1611.07308.pdf

compile(optimizer, loss=None, *args, **kwargs)[source]

Configures the model for training.

Parameters

optimizer (tf.keras.optimizers.Optimizer) – The optimizer to use for training.
loss (None, tf.keras.losses.Loss) – The loss function to use. If None, a default loss function will be used. This loss function simply just tried to maximize the values of the positive edges and minimize the values of the negative edges. If a custom loss function is used, be aware that the inputs to the loss functions are: positive edge scores (y_true ) and negative edge scores (y_pred); both resulting from molgraph.layers.DotProductIncident(…). Default to None.
metrics (None) – This argument will be ignored (at least for now).
*args – See tf.keras.Model.compile documentation.
**kwargs – See tf.keras.Model.compile documentation.

fit(x, y=None, batch_size=None, epochs=1, *args, **kwargs)[source]

Trains the autoencoder for a fixed number of epochs.

Parameters

x (GraphTensor, tf.data.Dataset) – The input data. Either a GraphTensor instance or a tf.data.Dataset of `GraphTensor`s.
y (None) – Target data, which will be ignored as the target can be obtained from x.
batch_size (int, None) – Number of samples per gradient update. If None, 32 will be used. Default to None.
epochs (int) – Number of iterations over all subgraphs (molecular graphs) of the GraphTensor instance or tf.data.Dataset instance. Default to 1.
*args – See tf.keras.Model.fit documentaton.
**kwargs – See tf.keras.Model.fit documentation.

Returns

A History object containing e.g. training loss values and metric values.

evaluate(x, y=None, batch_size=None, *args, **kwargs)[source]

Evaluates the autoencoder.

Parameters

x (GraphTensor, tf.data.Dataset) – The input data; either a GraphTensor instance or a tf.data.Dataset of `GraphTensor`s.
y (None) – Target data; which will be ignored as the target can be obtained from x.
batch_size (int, None) – Number of samples per batch of computation. If None, 32 will be used. Default to None.
*args – See tf.keras.Model.evaluate documentaton.
**kwargs – See tf.keras.Model.evaluate documentation.

Returns

Average loss values (e.g. reconstruction loss and kl loss).

predict(x, batch_size=None, *args, **kwargs)[source]

Generates outputs of the autoencoder.

Parameters

x (GraphTensor, tf.data.Dataset) – The input data; either a GraphTensor instance or a tf.data.Dataset of `GraphTensor`s.
batch_size (int, None) – Number of samples per batch of computation. If None, 32 will be used. Default to None.
*args – See tf.keras.Model.evaluate documentaton.
**kwargs – See tf.keras.Model.evaluate documentation.

Returns

tf.Tensor or tf.RaggedTensor of edge scores, corresponding to the inputted GraphTensor instance.

get_config()[source]

Returns the config of the Model.

Config is a Python dictionary (serializable) containing the configuration of an object, which in this case is a Model. This allows the Model to be be reinstantiated later (without its trained weights) from this configuration.

Note that get_config() does not guarantee to return a fresh copy of dict every time it is called. The callers should make a copy of the returned dict if they want to modify it.

Developers of subclassed Model are advised to override this method, and continue to update the dict from super(MyModel, self).get_config() to provide the proper configuration of this Model. The default config will return config dict for init parameters if they are basic types. Raises NotImplementedError when in cases where a custom get_config() implementation is required for the subclassed model.

Returns: Python dictionary containing the configuration of this Model.

classmethod from_config(config, custom_objects=None)[source]

Creates a layer from its config.

This method is the reverse of get_config, capable of instantiating the same layer from the config dictionary. It does not handle layer connectivity (handled by Network), nor weights (handled by set_weights).

Parameters: config – A Python dictionary, typically the output of get_config.
Returns: A layer instance.

GraphVariationalAutoEncoder

class molgraph.models.GraphVariationalAutoEncoder(tensorflow.keras.Model)[source]

Graph Variational AutoEncoder (GAE) based on Kipf and Welling 9.

Parameters

encoder (tf.keras.layers.Layer) – The encoder part of the autoencoder. The encoder could be any of the graph neural neural network layers provided by molgraph; e.g., GCNConv, GINConv or GATv2Conv.
decoder (tf.keras.layers.Layer) – The decoder part of the autoencoder. If None is passed, the decoder used is DotProductIncident. Default to None.
negative_graph_sampler (callable) – A function which samples negative graphs. It takes as input the current GraphTensor instance and produces a new GraphTensor instance with negative edges. If None, NegativeGraphSampler will be used. Default to None.
balanced_class_weighting (bool) – Whether balanced class weighting should be performed. For instance, if there are twice as many negative edges, should each negative edge be weighted 0.5 for the loss?
beta_initial (float) – Initial beta value (which is multiplied with the kl loss). Default to 0.0.
beta_end (float) – End beta value. Default to 0.1.
beta_incr (float) – The increment rate of the beta value (updated each train step). Default to 1e-6.

Example usage:

>>> # Replace this graph_tensor with a large dataset of graphs
>>> # Also accepts featurized graph tensors (via `chemistry.Featurizer`)
>>> graph_tensor = molgraph.GraphTensor(
...     sizes=[5],
...     edge_src=[1, 4, 0, 2, 3, 1, 1, 0],
...     edge_dst=[0, 0, 1, 1, 1, 2, 3, 4],
...     node_feature=['Sym:C|Hyb:SP3', 'Sym:C|Hyb:SP2', 'Sym:O|Hyb:SP2',
...                   'Sym:O|Hyb:SP2', 'Sym:N|Hyb:SP3'],
...     edge_feature=['BonTyp:SINGLE|Rot:1', 'BonTyp:SINGLE|Rot:0',
...                   'BonTyp:SINGLE|Rot:1', 'BonTyp:DOUBLE|Rot:0',
...                   'BonTyp:SINGLE|Rot:0', 'BonTyp:DOUBLE|Rot:0',
...                   'BonTyp:SINGLE|Rot:0', 'BonTyp:SINGLE|Rot:0'],
... )
>>> graph_tensor = graph_tensor.separate()
>>> node_embedding = molgraph.layers.NodeEmbeddingLookup(
...    32, mask_token='[MASK]')
>>> edge_embedding = molgraph.layers.EdgeEmbeddingLookup(
...    32, mask_token='[MASK]')
>>> node_embedding.adapt(graph_tensor)
>>> edge_embedding.adapt(graph_tensor)
>>> # Obtain the encoder of GVAE
>>> encoder_inputs = tf.keras.layers.Input(type_spec=graph_tensor.spec)
>>> encoder_x = node_embedding(encoder_inputs)
>>> encoder_x = edge_embedding(encoder_x)
>>> encoder_x = molgraph.layers.GATv2Conv(128, name='shared_conv')(encoder_x)
>>> encoder_x_mean = molgraph.layers.GATv2Conv(128, name='loc_conv')(encoder_x)
>>> encoder_x_log_var = molgraph.layers.GATv2Conv(128, name='log_var_conv')(encoder_x)
>>> encoder = tf.keras.Model(encoder_inputs, [encoder_x_mean, encoder_x_log_var])
>>> # Obtain the decoder of GVAE
>>> decoder = molgraph.layers.DotProductIncident(normalize=True)
>>> # Obtain, train and evaluate GVAE model
>>> gvae = molgraph.models.GraphVariationalAutoEncoder(encoder, decoder)
>>> gvae.compile('adam')
>>> _ = gvae.fit(graph_tensor, batch_size=32, epochs=50, verbose=0)
>>> total_loss, rec_loss, kl_loss = gvae.evaluate(graph_tensor, verbose=0)
>>> encoder.save( 
...     '/tmp/my_pretrained_encoder_model'
... )
>>> loaded_encoder = tf.saved_model.load( 
...     '/tmp/my_pretrained_encoder_model'
... )

References

9: https://arxiv.org/pdf/1611.07308.pdf

get_config()[source]

Returns the config of the Model.

Config is a Python dictionary (serializable) containing the configuration of an object, which in this case is a Model. This allows the Model to be be reinstantiated later (without its trained weights) from this configuration.

Note that get_config() does not guarantee to return a fresh copy of dict every time it is called. The callers should make a copy of the returned dict if they want to modify it.

Developers of subclassed Model are advised to override this method, and continue to update the dict from super(MyModel, self).get_config() to provide the proper configuration of this Model. The default config will return config dict for init parameters if they are basic types. Raises NotImplementedError when in cases where a custom get_config() implementation is required for the subclassed model.

Returns: Python dictionary containing the configuration of this Model.

fit(x, y=None, batch_size=None, epochs=1, *args, **kwargs)[source]

Trains the autoencoder for a fixed number of epochs.

Parameters

x (GraphTensor, tf.data.Dataset) – The input data. Either a GraphTensor instance or a tf.data.Dataset of `GraphTensor`s.
y (None) – Target data, which will be ignored as the target can be obtained from x.
batch_size (int, None) – Number of samples per gradient update. If None, 32 will be used. Default to None.
epochs (int) – Number of iterations over all subgraphs (molecular graphs) of the GraphTensor instance or tf.data.Dataset instance. Default to 1.
*args – See tf.keras.Model.fit documentaton.
**kwargs – See tf.keras.Model.fit documentation.

Returns

A History object containing e.g. training loss values and metric values.

evaluate(x, y=None, batch_size=None, *args, **kwargs)[source]

Evaluates the autoencoder.

Parameters

x (GraphTensor, tf.data.Dataset) – The input data; either a GraphTensor instance or a tf.data.Dataset of `GraphTensor`s.
y (None) – Target data; which will be ignored as the target can be obtained from x.
batch_size (int, None) – Number of samples per batch of computation. If None, 32 will be used. Default to None.
*args – See tf.keras.Model.evaluate documentaton.
**kwargs – See tf.keras.Model.evaluate documentation.

Returns

Average loss values (e.g. reconstruction loss and kl loss).

predict(x, batch_size=None, *args, **kwargs)[source]

Generates outputs of the autoencoder.

Parameters

x (GraphTensor, tf.data.Dataset) – The input data; either a GraphTensor instance or a tf.data.Dataset of `GraphTensor`s.
batch_size (int, None) – Number of samples per batch of computation. If None, 32 will be used. Default to None.
*args – See tf.keras.Model.evaluate documentaton.
**kwargs – See tf.keras.Model.evaluate documentation.

Returns

tf.Tensor or tf.RaggedTensor of edge scores, corresponding to the inputted GraphTensor instance.

compile(optimizer, loss=None, *args, **kwargs)[source]

Configures the model for training.

Parameters

optimizer (tf.keras.optimizers.Optimizer) – The optimizer to use for training.
loss (None, tf.keras.losses.Loss) – The loss function to use. If None, a default loss function will be used. This loss function simply just tried to maximize the values of the positive edges and minimize the values of the negative edges. If a custom loss function is used, be aware that the inputs to the loss functions are: positive edge scores (y_true ) and negative edge scores (y_pred); both resulting from molgraph.layers.DotProductIncident(…). Default to None.
metrics (None) – This argument will be ignored (at least for now).
*args – See tf.keras.Model.compile documentation.
**kwargs – See tf.keras.Model.compile documentation.

classmethod from_config(config, custom_objects=None)[source]

Creates a layer from its config.

This method is the reverse of get_config, capable of instantiating the same layer from the config dictionary. It does not handle layer connectivity (handled by Network), nor weights (handled by set_weights).

Parameters: config – A Python dictionary, typically the output of get_config.
Returns: A layer instance.