Layers

GNN Layer

class molgraph.layers.GNNLayer(tf.keras.layers.Layer)[source]

Base layer for the built-in GNN layers.

Can also be used to create new GNN layers.

Example usage:

>>> class MyGCNLayer(molgraph.layers.GNNLayer):
...
...     def __init__(self, units, **kwargs):
...         super().__init__(
...             units=units,
...             **kwargs)
...
...     def _call(self, graph_tensor):
...         node_feature_transformed = self.projection(
...             graph_tensor.node_feature)
...         graph_tensor = graph_tensor.update({
...             'node_feature': node_feature_transformed})
...         return graph_tensor.propagate()
...
...     def _build(self, graph_tensor_spec):
...         self.projection = self.get_dense(self.units)
...
>>> my_gcn_layer = MyGCNLayer(32)
>>> my_gcn_layer.compute_output_shape(tf.TensorShape((None, 8)))
TensorShape([None, 32])

Parameters

units (int, None) – Number of output units.
normalization – (None, str, bool): Whether to apply layer normalization to the output. If batch normalization is desired, pass ‘batch_norm’. Default to None.
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 output of the layer. 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. Default to tf.keras.initializers.TruncatedNormal(stddev=0.005).
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.
**kwargs –
Valid (optional) keyword arguments are:
- name (str): Name of the layer instance.
- update_step (tf.keras.layers.Layer): Applies post-processing step on the output (produced by _call). If passed, normalization, residual, activation and dropout parameters will be ignored. If None, a default post-processing step will be used (taking into consideration the aforementioned parameters). Default to None.
- use_edge_features: Whether or not to use edge features. Only relevant if edge features exist. If None, and edge features exist, it will be set to True. Default to None.
- update_edge_features (bool): Specifies whether edge features should be updated along with node features, including the post-processing step. Only relevant if edge features exist. It is important that GNN layers which updates its edge features for the next layer sets this to True. Default to False.

abstract _build(graph_tensor_spec)[source]

Builds the derived GNN layer.

Wrapped in build_from_signature().

Parameters: graph_tensor_spec (GraphTensor.Spec) – The spec corresponding to a graph tensor instance.

abstract _call(graph_tensor)[source]

Calls the derived GNN layer.

Wrapped in call().

Parameters: graph_tensor (GraphTensor) – Input to the layer.
Returns: A graph tensor with updated features. For some layers, only the node_feature data is updated, for other layers both the node_feature and the edge_feature data are updated.
Return type: GraphTensor

build_from_signature(graph_tensor)[source]

Builds the layer based on a GraphTensor or a Spec.

Automatically invoked on first call().

Parameters: graph_tensor (GraphTensor, GraphTensor.Spec) – A graph tensor instance or a spec of a graph tensor instance.

call(graph_tensor)[source]

Wraps _call(), applying pre- and post-processing steps.

Automatically builds the layer on first call.

Parameters: graph_tensor (GraphTensor) – An graph tensor instance.
Returns: An updated graph tensor instance.
Return type: GraphTensor

compute_output_shape(input_shape)[source]

Computes the output shape of the layer based on an input shape.

Parameters: input_shape (tf.TensorShape) – The shape of a GraphTensor instance.
Returns: The shape corresponding to the outputted (updated) GraphTensor instance.
Return type: tf.TensorShape

compute_output_signature(input_signature)[source]

Computes the output signature of the layer based on an input signature.

Parameters: input_signature (GraphTensor.Spec) – The spec of a GraphTensor instance.
Returns: The spec corrsponding to the outputted (updated) GraphTensor instance.
Return type: GraphTensor.Spec

get_bias(shape, dtype=tf.float32, name='bias')[source]

Obtain (trainable) bias weights based on parameters passed to __init__().

Parameters

shape (tf.TensorShape) – The shape of the weights.
dtype (tf.dtypes.DType) – The dtype of the weights. Default to tf.float32.
name (str) – The name of the weights. Default to ‘bias’.

Returns

Trainable bias weights.

Return type

tf.Variable

get_dense(units, activation=None, use_bias=None)[source]

Obtain a Dense layer based on parameters passed to __init__().

Parameters

units (int) – Number of units of the dense layer.
activation (None, callable) – The activation to be applied to the output of layer. Default to None.
use_bias (None, bool) – Whether bias should be used. If None, use_bias parameter passed to __init__ will be used. Default to None.

Returns

The dense layer.

Return type

layers.Dense

get_einsum_dense(equation, output_shape, activation=None, bias_axes=None)[source]

Obtain an EinsumDense layer based on parameters passed to __init__().

Parameters

equation (str) – The einsum formula.
output_shape (tf.TensorShape) – The output shape (excluding batch dimension or any dimensions represented by ellipses).
activation (None, callable) – The activation to be applied to the output of the layer. Default to None.
bias_axes (None, str) – A string containing the output dimension(s) to apply a bias to. The string characters should correspond to equation. If None and __init__(..., use_bias=True), the bias axes will be automatically derived. If None and __init__(..., use_bias=False), no bias will be applied.

Returns

The einsum dense layer.

Return type

layers.EinsumDense

get_kernel(shape, dtype=tf.float32, name='kernel')[source]

Obtain (trainable) kernel weights based on parameters passed to __init__().

Parameters

shape (tf.TensorShape) – The shape of the weights.
dtype (tf.dtypes.DType) – The dtype of the weights. Default to tf.float32.
name (str) – The name of the weights. Default to ‘kernel’.

Returns

Trainable bias weights.

Return type

tf.Variable

GNN ops

GNN ops are helper functions which makes it easier to code up a custom GNN layer. For example, a basic GCN layer can be coded up as follows:

import tensorflow as tf
from molgraph.layers import gnn_ops

class MyGCNConv(tf.keras.layers.Layer):

  def __init__(self, units):
      super().__init__()
      self.units = units

  def build(self, input_shape):
      self.kernel = self.add_weight(
          name='kernel',
          shape=(input_shape[-1], self.units),
          dtype=tf.float32,
          trainable=True)
      self.built = True

  def call(self, graph_tensor):
      graph_tensor_orig = graph_tensor
      if isinstance(graph_tensor.node_feature, tf.RaggedTensor):
          graph_tensor = graph_tensor.merge()
      node_feature_transformed = tf.matmul(graph_tensor.node_feature, self.kernel)
      node_feature_aggregated = gnn_ops.propagate_node_features(
          node_feature=node_feature_transformed,
          edge_dst=graph_tensor.edge_dst,
          edge_src=graph_tensor.edge_src,
          mode='mean')
      return graph_tensor_orig.update({'node_feature': node_feature_aggregated})

molgraph.layers.gnn_ops.propagate_node_features(*, node_feature, edge_src, edge_dst, edge_weight=None, mode='sum')[source]

Aggregates source node features to destination node features.

Parameters

node_feature (tf.Tensor) – Node features of GraphTensor instance.
edge_src (tf.Tensor) – Source node indices of GraphTensor instance. Entry i corresponds to row i in node_feature.
edge_dst (tf.Tensor) – Destination node indices of GraphTensor instance. Entry i corresponds to row i in node_feature.
edge_weight (tf.Tensor) – Edge weights associated with edge_dst and edge_src.
mode (str) – The type of aggregation to be performed, either of ‘sum’, ‘mean’, ‘min’ or ‘max’. If None, ‘sum’ will be used. Default to ‘sum’.

Returns

Updated node features.

Return type

tf.Tensor

molgraph.layers.gnn_ops.softmax_edge_weights(*, edge_weight, edge_dst, exponentiate=True, clip_values=(-5.0, 5.0))[source]

Normalizes edge weights via softmax.

The sum of edge weights associated with each destination node sums to 1.

Parameters

edge_weight (tf.Tensor) – Edge weights associated with edge_src and edge_dst of the GraphTensor instance.
edge_dst (tf.Tensor) – Destination node indices (corresponding to edge_weight) of the GraphTensor instance.
exponentiate (bool) – Whether edge_weight should be exponentiated. Default to True.
clip_values (tuple) – If exponentiation, clip values before it, for stability.

Returns

Normalized edge weights.

Return type

tf.Tensor

molgraph.layers.gnn_ops.compute_edge_weights_from_degrees(*, edge_src, edge_dst, edge_feature=None, mode='symmetric')[source]

Computes edge weights based on the number of nearest neighbors.

These edge weights can subsequently be used as normalization coefficients for the neighborhood aggregation (propagate_node_features).

Parameters

edge_src (tf.Tensor) – Source node indices.
edge_dst (tf.Tensor) – Destination node indices.
edge_feature (tf.Tensor) – Edge features associated with edge_dst and edge_src.
mode (str) – Type of normalization to use. Either of ‘symmetric’, ‘row’ or None. If None, ‘row’ is used. Default to ‘symmetric’.

Returns

Edge weights based on degrees.

Return type

tf.Tensor

molgraph.layers.gnn_ops.reduce_features(*, feature, mode, output_units=None, axis=1)[source]

Reduces dimension of features.

Useful to concatenate, sum or average heads of GATConv or GTConv.

Parameters

feature (tf.Tensor) – The features to be reduced. Either node or edge features.
mode (str, None) – The type of reduction to be performed. Either of ‘concat’, ‘sum’, ‘mean’ or None. If None, ‘mean’ is performed.
output_units (int, None) – The output dimension (innermost dimension) after reshaping. Only relevant if mode='concat'. If None, dim is inferred from the two innermost dimensions. Default to None,
axis (int) – Axis to reduce. Default to 1.

Returns

Reduced features (rank 3 -> rank 2).

Return type

tf.Tensor

Convolutional

GCNConv

class molgraph.layers.GCNConv(molgraph.layers.GNNLayer)[source]

Graph convolutional layer (GCN).

Implementation is based on Kipf et al. (2017) 1, Dwivedi et al. (2022) 2, and, for RGCN, Schlichtkrull et al. (2017) 3.

Example usage:

>>> graph_tensor = molgraph.GraphTensor(
...     sizes=[2, 3],
...     node_feature=[[1., 0.], [1., 0.], [1., 0.], [1., 0.], [0., 1.]],
...     edge_src=[1, 0, 3, 4, 2, 4, 3, 2],
...     edge_dst=[0, 1, 2, 2, 3, 3, 4, 4],
... )
>>> gnn_model = tf.keras.Sequential([
...     molgraph.layers.GCNConv(units=16),
...     molgraph.layers.GCNConv(units=16),
...     molgraph.layers.GCNConv(units=16),
...     molgraph.layers.Readout(),
... ])
>>> gnn_model(graph_tensor).shape
TensorShape([2, 16])

Parameters

units (int, None) – Number of output units.
use_edge_features (bool) – Whether or not to use edge features. Default to False.
degree_normalization (str, None) – The strategy for computing edge weights from degrees. Either of ‘symmetric’, ‘row’ or None. If None, ‘row’ is used. Default to ‘symmetric’.
num_bases (int, None) – Number of bases to use for basis decomposition. Only relevant if use_edge_features is True. Default to None.
self_projection (bool) – Whether to apply self projection. Default to True.
normalization – (None, str, bool): Whether to apply layer normalization to the output. If batch normalization is desired, pass ‘batch_norm’. Default to None.
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 output of the layer. 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. Default to tf.keras.initializers.TruncatedNormal(stddev=0.005).
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.
**kwargs –
Valid (optional) keyword arguments are:
- name (str): Name of the layer instance.
- update_step (tf.keras.layers.Layer): Applies post-processing step on the output (produced by _call). If passed, normalization, residual, activation and dropout parameters will be ignored. If None, a default post-processing step will be used (taking into consideration the aforementioned parameters). Default to None.
- use_edge_features: Whether or not to use edge features. Only relevant if edge features exist. Default to False.

References

1: https://arxiv.org/pdf/1609.02907.pdf
2: https://arxiv.org/pdf/2003.00982.pdf
3: https://arxiv.org/pdf/1703.06103.pdf

call(graph_tensor)[source]

Wraps _call(), applying pre- and post-processing steps.

Automatically builds the layer on first call.

Parameters: graph_tensor (GraphTensor) – An graph tensor instance.
Returns: An updated graph tensor instance.
Return type: GraphTensor

compute_output_shape(input_shape)[source]

Computes the output shape of the layer based on an input shape.

Parameters: input_shape (tf.TensorShape) – The shape of a GraphTensor instance.
Returns: The shape corresponding to the outputted (updated) GraphTensor instance.
Return type: tf.TensorShape

classmethod from_config(config)[source]

Initializes and builds the layer based on a configuration.

Parameters: config (dict) – A Python dictionary of parameters to initialize and build a layer instance. The config is usually obtained from get_config() of another layer instance (of the same class).
Returns: A layer instance.

get_config()[source]

Returns the configuration of the layer.

Returns: Python dictionary with the layer’s parameters. Can be used to initialize and build another layer instance (of the same class), via from_config().

GINConv

class molgraph.layers.GINConv(molgraph.layers.GNNLayer)[source]

Graph isomorphism convolution layer (GIN).

Implementation based on Dwivedi et al. (2022) 4, Xu et al. (2019) 5, and Hu et al. (2020) 6.

Example usage:

>>> graph_tensor = molgraph.GraphTensor(
...     sizes=[2, 3],
...     node_feature=[[1., 0.], [1., 0.], [1., 0.], [1., 0.], [0., 1.]],
...     edge_src=[1, 0, 3, 4, 2, 4, 3, 2],
...     edge_dst=[0, 1, 2, 2, 3, 3, 4, 4],
... )
>>> gnn_model = tf.keras.Sequential([
...     molgraph.layers.GINConv(units=16),
...     molgraph.layers.GINConv(units=16),
...     molgraph.layers.GINConv(units=16),
...     molgraph.layers.Readout(),
... ])
>>> gnn_model(graph_tensor).shape
TensorShape([2, 16])

Including edge features:

>>> 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],
... )
>>> gnn_model = tf.keras.Sequential([
...     molgraph.layers.GINConv(units=16, use_edge_features=True),
...     molgraph.layers.GINConv(units=16, use_edge_features=True),
...     molgraph.layers.GINConv(units=16, use_edge_features=True),
... ])
>>> output = gnn_model(graph_tensor)
>>> output.node_feature.shape, output.edge_feature.shape
(TensorShape([5, 16]), TensorShape([8, 16]))

Parameters

units (int, None) – Number of output units.
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.
self_projection (bool) – Whether to apply self projection. Default to True.
normalization – (None, str, bool): Whether to apply layer normalization to the output. If batch normalization is desired, pass ‘batch_norm’. Default to None.
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 output of the layer. 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. Default to tf.keras.initializers.TruncatedNormal(stddev=0.005).
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.
**kwargs –
Valid (optional) keyword arguments are:
- name (str): Name of the layer instance.
- update_step (tf.keras.layers.Layer): Applies post-processing step on the output (produced by _call). If passed, normalization, residual, activation and dropout parameters will be ignored. If None, a default post-processing step will be used (taking into consideration the aforementioned parameters). Default to None.
- use_edge_features: Whether or not to use edge features. Only relevant if edge features exist. If None, and edge features exist, it will be set to True. Default to None.

References

4: https://arxiv.org/pdf/2003.00982.pdf
5: https://arxiv.org/pdf/1810.00826.pdf
6: https://arxiv.org/pdf/1905.12265.pdf

call(graph_tensor)[source]

Wraps _call(), applying pre- and post-processing steps.

Automatically builds the layer on first call.

Parameters: graph_tensor (GraphTensor) – An graph tensor instance.
Returns: An updated graph tensor instance.
Return type: GraphTensor

compute_output_shape(input_shape)[source]

Computes the output shape of the layer based on an input shape.

Parameters: input_shape (tf.TensorShape) – The shape of a GraphTensor instance.
Returns: The shape corresponding to the outputted (updated) GraphTensor instance.
Return type: tf.TensorShape

classmethod from_config(config)[source]

Initializes and builds the layer based on a configuration.

Parameters: config (dict) – A Python dictionary of parameters to initialize and build a layer instance. The config is usually obtained from get_config() of another layer instance (of the same class).
Returns: A layer instance.

get_config()[source]

Returns the configuration of the layer.

Returns: Python dictionary with the layer’s parameters. Can be used to initialize and build another layer instance (of the same class), via from_config().

GraphSageConv

class molgraph.layers.GraphSageConv(molgraph.layers.GNNLayer)[source]

Graph sage convolution layer (GraphSage)

Implementation is based on Hamilton et al. (2018) 7 and Dwivedi et al. (2022) 8.

Example usage:

>>> graph_tensor = molgraph.GraphTensor(
...     sizes=[2, 3],
...     node_feature=[[1., 0.], [1., 0.], [1., 0.], [1., 0.], [0., 1.]],
...     edge_src=[1, 0, 3, 4, 2, 4, 3, 2],
...     edge_dst=[0, 1, 2, 2, 3, 3, 4, 4],
... )
>>> gnn_model = tf.keras.Sequential([
...     molgraph.layers.GraphSageConv(units=16),
...     molgraph.layers.GraphSageConv(units=16),
...     molgraph.layers.GraphSageConv(units=16),
...     molgraph.layers.Readout(),
... ])
>>> gnn_model(graph_tensor).shape
TensorShape([2, 16])

Parameters

units (int, None) – Number of output units.
aggregation_mode (str) – Type of neighborhood aggregation to be performed. Either of ‘max’, ‘lstm’, ‘mean’, ‘sum’. Default to ‘mean’.
normalize (bool) – Whether l2 normalization should be performed on the updated node features. Default to True.
self_projection (bool) – Whether to apply self projection. Default to True.
normalization – (None, str, bool): Whether to apply layer normalization to the output. If batch normalization is desired, pass ‘batch_norm’. Default to None.
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 output of the layer. 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. Default to tf.keras.initializers.TruncatedNormal(stddev=0.005).
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.
**kwargs –
Valid (optional) keyword arguments are:
- name (str): Name of the layer instance.
- update_step (tf.keras.layers.Layer): Applies post-processing step on the output (produced by _call). If passed, normalization, residual, activation and dropout parameters will be ignored. If None, a default post-processing step will be used (taking into consideration the aforementioned parameters). Default to None.

References

7: https://arxiv.org/pdf/1706.02216.pdf
8: https://arxiv.org/pdf/2003.00982.pdf

call(graph_tensor)[source]

Wraps _call(), applying pre- and post-processing steps.

Automatically builds the layer on first call.

Parameters: graph_tensor (GraphTensor) – An graph tensor instance.
Returns: An updated graph tensor instance.
Return type: GraphTensor

compute_output_shape(input_shape)[source]

Computes the output shape of the layer based on an input shape.

Parameters: input_shape (tf.TensorShape) – The shape of a GraphTensor instance.
Returns: The shape corresponding to the outputted (updated) GraphTensor instance.
Return type: tf.TensorShape

classmethod from_config(config)[source]

Initializes and builds the layer based on a configuration.

Parameters: config (dict) – A Python dictionary of parameters to initialize and build a layer instance. The config is usually obtained from get_config() of another layer instance (of the same class).
Returns: A layer instance.

get_config()[source]

Returns the configuration of the layer.

Returns: Python dictionary with the layer’s parameters. Can be used to initialize and build another layer instance (of the same class), via from_config().

GCNIIConv

class molgraph.layers.GCNIIConv(molgraph.layers.GNNLayer)[source]

Graph convolutional ‘via Initial residual and Identity mapping’ layer (GCNII).

Implementation is based on Chen et al. (2020) 9.

Example usage:

>>> graph_tensor = molgraph.GraphTensor(
...     sizes=[2, 3],
...     node_feature=[[1., 0.], [1., 0.], [1., 0.], [1., 0.], [0., 1.]],
...     edge_src=[1, 0, 3, 4, 2, 4, 3, 2],
...     edge_dst=[0, 1, 2, 2, 3, 3, 4, 4],
... )
>>> gnn_model = tf.keras.Sequential([
...     molgraph.layers.GCNIIConv(units=16),
...     molgraph.layers.GCNIIConv(units=16),
...     molgraph.layers.GCNIIConv(units=16),
...     molgraph.layers.Readout(),
... ])
>>> gnn_model(graph_tensor).shape
TensorShape([2, 16])

Parameters

units (int, None) – Number of output units.
alpha (float) – Decides how much information of the initial node state (the original node features) should be passed to the next layer (alpha) vs. how much new information should be passed to the subsequent layer (1 - alpha). Takes a value between 0.0 and 1.0. In the original paper, alpha was set between 0.1 and 0.5, depending on the dataset.
beta (float) – Decides to what extent the kernel (projection) should be ignored. Takes a value between 0.0 and 1.0; a value set to 0.0 means that the kernel is ignored, a value set to 1.0 means that the kernel is fully applied. In the original paper, beta is set to log(lambda/l+1); l denotes the l:th layer, and lambda is a hyperparameter, set, in the original paper, between 0.5 and 1.5 depending on the dataset.
variant (bool) – Whether the GCNII variant should be used. Default to False.
degree_normalization (str, None) – The strategy for computing edge weights from degrees. Either of ‘symmetric’, ‘row’ or None. If None, ‘row’ is used. Default to ‘symmetric’.
self_projection (bool) – Whether to apply self projection. Default to True.
normalization – (None, str, bool): Whether to apply layer normalization to the output. If batch normalization is desired, pass ‘batch_norm’. Default to None.
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 output of the layer. 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. Default to tf.keras.initializers.TruncatedNormal(stddev=0.005).
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.
**kwargs –
Valid (optional) keyword arguments are:
- name (str): Name of the layer instance.
- update_step (tf.keras.layers.Layer): Applies post-processing step on the output (produced by _call). If passed, normalization, residual, activation and dropout parameters will be ignored. If None, a default post-processing step will be used (taking into consideration the aforementioned parameters). Default to None.

References

9: https://arxiv.org/pdf/2007.02133v1.pdf

call(graph_tensor)[source]

Wraps _call(), applying pre- and post-processing steps.

Automatically builds the layer on first call.

Parameters: graph_tensor (GraphTensor) – An graph tensor instance.
Returns: An updated graph tensor instance.
Return type: GraphTensor

compute_output_shape(input_shape)[source]

Computes the output shape of the layer based on an input shape.

Parameters: input_shape (tf.TensorShape) – The shape of a GraphTensor instance.
Returns: The shape corresponding to the outputted (updated) GraphTensor instance.
Return type: tf.TensorShape

classmethod from_config(config)[source]

Initializes and builds the layer based on a configuration.

Parameters: config (dict) – A Python dictionary of parameters to initialize and build a layer instance. The config is usually obtained from get_config() of another layer instance (of the same class).
Returns: A layer instance.

get_config()[source]

Returns the configuration of the layer.

Returns: Python dictionary with the layer’s parameters. Can be used to initialize and build another layer instance (of the same class), via from_config().

Attentional

GATConv

class molgraph.layers.GATConv(molgraph.layers.GNNLayer)[source]

Multi-head graph attention layer (GAT).

The implementation is based on Velickovic et al. (2018) 10 and Dwivedi et al. (2022) 11.

Example usage:

>>> graph_tensor = molgraph.GraphTensor(
...     sizes=[2, 3],
...     node_feature=[[1., 0.], [1., 0.], [1., 0.], [1., 0.], [0., 1.]],
...     edge_src=[1, 0, 3, 4, 2, 4, 3, 2],
...     edge_dst=[0, 1, 2, 2, 3, 3, 4, 4],
... )
>>> gnn_model = tf.keras.Sequential([
...     molgraph.layers.GATConv(units=16),
...     molgraph.layers.GATConv(units=16),
...     molgraph.layers.GATConv(units=16),
...     molgraph.layers.Readout(),
... ])
>>> gnn_model(graph_tensor).shape
TensorShape([2, 16])

Including edge features:

>>> 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],
... )
>>> gnn_model = tf.keras.Sequential([
...     molgraph.layers.GATConv(units=16, use_edge_features=True),
...     molgraph.layers.GATConv(units=16, use_edge_features=True),
...     molgraph.layers.GATConv(units=16, use_edge_features=True),
... ])
>>> output = gnn_model(graph_tensor)
>>> output.node_feature.shape, output.edge_feature.shape
(TensorShape([5, 16]), TensorShape([8, 16]))

Parameters

units (int, None) – Number of output units.
num_heads (int) – Number of attention heads. Default to 8.
merge_mode (str) – The strategy for merging the heads. Either of ‘concat’, ‘sum’, ‘mean’ or None. If set to None, ‘mean’ is used. Default to ‘concat’.
self_projection (bool) – Whether to apply self projection. Default to True.
normalization – (None, str, bool): Whether to apply layer normalization to the output. If batch normalization is desired, pass ‘batch_norm’. Default to None.
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.
attention_activation (tf.keras.activations.Activation, callable, str, None) – Activation function applied to the the attention scores. Default to ‘leaky_relu’.
activation (tf.keras.activations.Activation, callable, str, None) – Activation function applied to the output of the layer. 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. Default to tf.keras.initializers.TruncatedNormal(stddev=0.005).
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.
**kwargs –
Valid (optional) keyword arguments are:
- name (str): Name of the layer instance.
- update_step (tf.keras.layers.Layer): Applies post-processing step on the output (produced by _call). If passed, normalization, residual, activation and dropout parameters will be ignored. If None, a default post-processing step will be used (taking into consideration the aforementioned parameters). Default to None.
- use_edge_features: Whether or not to use edge features. Only relevant if edge features exist. If None, and edge features exist, it will be set to True. Default to None.
- update_edge_features (bool): Specifies whether edge features should be updated along with node features, including the post-processing step. Only relevant if edge features exist. It is important that GNN layers which updates its edge features for the next layer sets this to True. Default to False.

References

10: https://arxiv.org/pdf/1710.10903.pdf
11: https://arxiv.org/pdf/2003.00982.pdf

call(graph_tensor)[source]

Wraps _call(), applying pre- and post-processing steps.

Automatically builds the layer on first call.

Parameters: graph_tensor (GraphTensor) – An graph tensor instance.
Returns: An updated graph tensor instance.
Return type: GraphTensor

compute_output_shape(input_shape)[source]

Computes the output shape of the layer based on an input shape.

Parameters: input_shape (tf.TensorShape) – The shape of a GraphTensor instance.
Returns: The shape corresponding to the outputted (updated) GraphTensor instance.
Return type: tf.TensorShape

classmethod from_config(config)[source]

Initializes and builds the layer based on a configuration.

Parameters: config (dict) – A Python dictionary of parameters to initialize and build a layer instance. The config is usually obtained from get_config() of another layer instance (of the same class).
Returns: A layer instance.

get_config()[source]

Returns the configuration of the layer.

Returns: Python dictionary with the layer’s parameters. Can be used to initialize and build another layer instance (of the same class), via from_config().

GATv2Conv

class molgraph.layers.GATv2Conv(molgraph.layers.GNNLayer)[source]

Multi-head graph attention (v2) layer (GATv2).

The implementation is based on Brody et al. (2021) 12.

Example usage:

>>> graph_tensor = molgraph.GraphTensor(
...     sizes=[2, 3],
...     node_feature=[[1., 0.], [1., 0.], [1., 0.], [1., 0.], [0., 1.]],
...     edge_src=[1, 0, 3, 4, 2, 4, 3, 2],
...     edge_dst=[0, 1, 2, 2, 3, 3, 4, 4],
... )
>>> gnn_model = tf.keras.Sequential([
...     molgraph.layers.GATv2Conv(units=16),
...     molgraph.layers.GATv2Conv(units=16),
...     molgraph.layers.GATv2Conv(units=16),
...     molgraph.layers.Readout(),
... ])
>>> gnn_model(graph_tensor).shape
TensorShape([2, 16])

Including edge features:

>>> 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],
... )
>>> gnn_model = tf.keras.Sequential([
...     molgraph.layers.GATv2Conv(units=16, use_edge_features=True),
...     molgraph.layers.GATv2Conv(units=16, use_edge_features=True),
...     molgraph.layers.GATv2Conv(units=16, use_edge_features=True),
... ])
>>> output = gnn_model(graph_tensor)
>>> output.node_feature.shape, output.edge_feature.shape
(TensorShape([5, 16]), TensorShape([8, 16]))

Parameters

units (int, None) – Number of output units.
num_heads (int) – Number of attention heads. Default to 8.
merge_mode (str) – The strategy for merging the heads. Either of ‘concat’, ‘sum’, ‘mean’ or None. If set to None, ‘mean’ is used. Default to ‘concat’.
self_projection (bool) – Whether to apply self projection. Default to True.
normalization – (None, str, bool): Whether to apply layer normalization to the output. If batch normalization is desired, pass ‘batch_norm’. Default to None.
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.
attention_activation (tf.keras.activations.Activation, callable, str, None) – Activation function applied to the the attention scores. Default to ‘leaky_relu’.
activation (tf.keras.activations.Activation, callable, str, None) – Activation function applied to the output of the layer. 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. Default to tf.keras.initializers.TruncatedNormal(stddev=0.005).
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.
**kwargs –
Valid (optional) keyword arguments are:
- name (str): Name of the layer instance.
- update_step (tf.keras.layers.Layer): Applies post-processing step on the output (produced by _call). If passed, normalization, residual, activation and dropout parameters will be ignored. If None, a default post-processing step will be used (taking into consideration the aforementioned parameters). Default to None.
- use_edge_features: Whether or not to use edge features. Only relevant if edge features exist. If None, and edge features exist, it will be set to True. Default to None.
- update_edge_features (bool): Specifies whether edge features should be updated along with node features, including the post-processing step. Only relevant if edge features exist. It is important that GNN layers which updates its edge features for the next layer sets this to True. Default to False.

References

12: https://arxiv.org/pdf/2105.14491.pdf

call(graph_tensor)[source]

Wraps _call(), applying pre- and post-processing steps.

Automatically builds the layer on first call.

Parameters: graph_tensor (GraphTensor) – An graph tensor instance.
Returns: An updated graph tensor instance.
Return type: GraphTensor

compute_output_shape(input_shape)[source]

Computes the output shape of the layer based on an input shape.

Parameters: input_shape (tf.TensorShape) – The shape of a GraphTensor instance.
Returns: The shape corresponding to the outputted (updated) GraphTensor instance.
Return type: tf.TensorShape

classmethod from_config(config)[source]

Initializes and builds the layer based on a configuration.

Parameters: config (dict) – A Python dictionary of parameters to initialize and build a layer instance. The config is usually obtained from get_config() of another layer instance (of the same class).
Returns: A layer instance.

get_config()[source]

Returns the configuration of the layer.

Returns: Python dictionary with the layer’s parameters. Can be used to initialize and build another layer instance (of the same class), via from_config().

GatedGCNConv

class molgraph.layers.GatedGCNConv(molgraph.layers.GNNLayer)[source]

Gated graph convolutional layer (GatedGCN).

Implementation is based on Dwivedi et al. (2022) 13, Bresson et al. (2019) 14, Joshi et al. (2019) 15, and Bresson et al. (2018) 16.

Example usage:

>>> graph_tensor = molgraph.GraphTensor(
...     sizes=[2, 3],
...     node_feature=[[1., 0.], [1., 0.], [1., 0.], [1., 0.], [0., 1.]],
...     edge_src=[1, 0, 3, 4, 2, 4, 3, 2],
...     edge_dst=[0, 1, 2, 2, 3, 3, 4, 4],
... )
>>> gnn_model = tf.keras.Sequential([
...     molgraph.layers.GatedGCNConv(units=16),
...     molgraph.layers.GatedGCNConv(units=16),
...     molgraph.layers.GatedGCNConv(units=16),
...     molgraph.layers.Readout(),
... ])
>>> gnn_model(graph_tensor).shape
TensorShape([2, 16])

Including edge features:

>>> 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],
... )
>>> gnn_model = tf.keras.Sequential([
...     molgraph.layers.GatedGCNConv(units=16, use_edge_features=True),
...     molgraph.layers.GatedGCNConv(units=16, use_edge_features=True),
...     molgraph.layers.GatedGCNConv(units=16, use_edge_features=True),
... ])
>>> output = gnn_model(graph_tensor)
>>> output.node_feature.shape, output.edge_feature.shape
(TensorShape([5, 16]), TensorShape([8, 16]))

Parameters

units (int, None) – Number of output units.
self_projection (bool) – Whether to apply self projection. Default to True.
normalization – (None, str, bool): Whether to apply layer normalization to the output. If batch normalization is desired, pass ‘batch_norm’. Default to None.
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 output of the layer. 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. Default to tf.keras.initializers.TruncatedNormal(stddev=0.005).
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.
**kwargs –
Valid (optional) keyword arguments are:
- name (str): Name of the layer instance.
- update_step (tf.keras.layers.Layer): Applies post-processing step on the output (produced by _call). If passed, normalization, residual, activation and dropout parameters will be ignored. If None, a default post-processing step will be used (taking into consideration the aforementioned parameters). Default to None.
- use_edge_features: Whether or not to use edge features. Only relevant if edge features exist. If None, and edge features exist, it will be set to True. Default to None.
- update_edge_features (bool): Specifies whether edge features should be updated along with node features, including the post-processing step. Only relevant if edge features exist. It is important that GNN layers which updates its edge features for the next layer sets this to True. Default to False.

References

13: https://arxiv.org/pdf/2003.00982.pdf
14: https://arxiv.org/pdf/1906.03412.pdf
15: https://arxiv.org/pdf/1906.01227.pdf
16: https://arxiv.org/pdf/1711.07553.pdf

call(graph_tensor)[source]

Wraps _call(), applying pre- and post-processing steps.

Automatically builds the layer on first call.

Parameters: graph_tensor (GraphTensor) – An graph tensor instance.
Returns: An updated graph tensor instance.
Return type: GraphTensor

compute_output_shape(input_shape)[source]

Computes the output shape of the layer based on an input shape.

Parameters: input_shape (tf.TensorShape) – The shape of a GraphTensor instance.
Returns: The shape corresponding to the outputted (updated) GraphTensor instance.
Return type: tf.TensorShape

classmethod from_config(config)[source]

Initializes and builds the layer based on a configuration.

Parameters: config (dict) – A Python dictionary of parameters to initialize and build a layer instance. The config is usually obtained from get_config() of another layer instance (of the same class).
Returns: A layer instance.

get_config()[source]

Returns the configuration of the layer.

Returns: Python dictionary with the layer’s parameters. Can be used to initialize and build another layer instance (of the same class), via from_config().

GMMConv

class molgraph.layers.GMMConv(molgraph.layers.GNNLayer)[source]

Multi-head graph gaussian mixture layer (MoNet)

Implementation is based on Dwivedi et al. (2022) 17 and Monti et al. (2016) 18.

Example usage:

>>> graph_tensor = molgraph.GraphTensor(
...     sizes=[2, 3],
...     node_feature=[[1., 0.], [1., 0.], [1., 0.], [1., 0.], [0., 1.]],
...     edge_src=[1, 0, 3, 4, 2, 4, 3, 2],
...     edge_dst=[0, 1, 2, 2, 3, 3, 4, 4],
... )
>>> gnn_model = tf.keras.Sequential([
...     molgraph.layers.GMMConv(units=16),
...     molgraph.layers.GMMConv(units=16),
...     molgraph.layers.GMMConv(units=16),
...     molgraph.layers.Readout(),
... ])
>>> gnn_model(graph_tensor).shape
TensorShape([2, 16])

Parameters

units (int, None) – Number of output units.
num_kernels (int) – Number of attention heads. Default to 8.
merge_mode (str) – The strategy for merging the heads. Either of ‘concat’, ‘sum’, ‘mean’ or None. If set to None, ‘mean’ is used. Default to ‘sum’.
pseudo_coord_dim (int) – The dimension of the pseudo coordinate of the Gaussian kernel. Default to 2.
self_projection (bool) – Whether to apply self projection. Default to True.
normalization – (None, str, bool): Whether to apply layer normalization to the output. If batch normalization is desired, pass ‘batch_norm’. Default to None.
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 output of the layer. 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. Default to tf.keras.initializers.TruncatedNormal(stddev=0.005).
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.
**kwargs –
Valid (optional) keyword arguments are:
- name (str): Name of the layer instance.
- update_step (tf.keras.layers.Layer): Applies post-processing step on the output (produced by _call). If passed, normalization, residual, activation and dropout parameters will be ignored. If None, a default post-processing step will be used (taking into consideration the aforementioned parameters). Default to None.

References:

17: https://arxiv.org/pdf/2003.00982.pdf
18: https://arxiv.org/pdf/1611.08402.pdf

call(graph_tensor)[source]

Wraps _call(), applying pre- and post-processing steps.

Automatically builds the layer on first call.

Parameters: graph_tensor (GraphTensor) – An graph tensor instance.
Returns: An updated graph tensor instance.
Return type: GraphTensor

compute_output_shape(input_shape)[source]

Computes the output shape of the layer based on an input shape.

Parameters: input_shape (tf.TensorShape) – The shape of a GraphTensor instance.
Returns: The shape corresponding to the outputted (updated) GraphTensor instance.
Return type: tf.TensorShape

classmethod from_config(config)[source]

Initializes and builds the layer based on a configuration.

Parameters: config (dict) – A Python dictionary of parameters to initialize and build a layer instance. The config is usually obtained from get_config() of another layer instance (of the same class).
Returns: A layer instance.

get_config()[source]

Returns the configuration of the layer.

Returns: Python dictionary with the layer’s parameters. Can be used to initialize and build another layer instance (of the same class), via from_config().

GTConv

class molgraph.layers.GTConv(molgraph.layers.GNNLayer)[source]

Graph transformer layer

Implementation is based on Dwivedi et al. (2021) 19.

Alias: GraphTransformerConv

Example usage:

>>> graph_tensor = molgraph.GraphTensor(
...     sizes=[2, 3],
...     node_feature=[[1., 0.], [1., 0.], [1., 0.], [1., 0.], [0., 1.]],
...     edge_src=[1, 0, 3, 4, 2, 4, 3, 2],
...     edge_dst=[0, 1, 2, 2, 3, 3, 4, 4],
... )
>>> gnn_model = tf.keras.Sequential([
...     molgraph.layers.GTConv(units=16),
...     molgraph.layers.GTConv(units=16),
...     molgraph.layers.GTConv(units=16),
...     molgraph.layers.Readout(),
... ])
>>> gnn_model(graph_tensor).shape
TensorShape([2, 16])

Including edge features:

>>> 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],
... )
>>> gnn_model = tf.keras.Sequential([
...     molgraph.layers.GTConv(units=16, use_edge_features=True),
...     molgraph.layers.GTConv(units=16, use_edge_features=True),
...     molgraph.layers.GTConv(units=16, use_edge_features=True),
... ])
>>> output = gnn_model(graph_tensor)
>>> output.node_feature.shape, output.edge_feature.shape
(TensorShape([5, 16]), TensorShape([8, 16]))

Parameters

units (int, None) – Number of output units.
num_heads (int) – Number of attention heads. Default to 8.
merge_mode (str) – The strategy for merging the heads. Either of ‘concat’, ‘sum’, ‘mean’ or None. If set to None, ‘mean’ is used. Default to ‘concat’.
self_projection (bool) – Whether to apply self projection. Default to True.
normalization (str, None) – The type of normalization to use for the output. Either of ‘batch_norm’, ‘layer_norm’ or None. Default to ‘layer_norm’.
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.
attention_activation (tf.keras.activations.Activation, callable, str, None) – Activation function applied to the the attention scores. Default to None.
activation (tf.keras.activations.Activation, callable, str, None) – Activation function applied to the output of the layer. 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. Default to tf.keras.initializers.TruncatedNormal(stddev=0.005).
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.
**kwargs –
Valid (optional) keyword arguments are:
- name (str): Name of the layer instance.
- update_step (tf.keras.layers.Layer): Applies post-processing step on the output (produced by _call). If passed, normalization, residual, activation and dropout parameters will be ignored. If None, a default post-processing step will be used (taking into consideration the aforementioned parameters). Default to None.
- use_edge_features: Whether or not to use edge features. Only relevant if edge features exist. If None, and edge features exist, it will be set to True. Default to None.
- update_edge_features (bool): Specifies whether edge features should be updated along with node features, including the post-processing step. Only relevant if edge features exist. It is important that GNN layers which updates its edge features for the next layer sets this to True. Default to False.

References

19: https://arxiv.org/pdf/2012.09699.pdf

call(graph_tensor)[source]

Wraps _call(), applying pre- and post-processing steps.

Automatically builds the layer on first call.

Parameters: graph_tensor (GraphTensor) – An graph tensor instance.
Returns: An updated graph tensor instance.
Return type: GraphTensor

compute_output_shape(input_shape)[source]

Computes the output shape of the layer based on an input shape.

Parameters: input_shape (tf.TensorShape) – The shape of a GraphTensor instance.
Returns: The shape corresponding to the outputted (updated) GraphTensor instance.
Return type: tf.TensorShape

classmethod from_config(config)[source]

Initializes and builds the layer based on a configuration.

Parameters: config (dict) – A Python dictionary of parameters to initialize and build a layer instance. The config is usually obtained from get_config() of another layer instance (of the same class).
Returns: A layer instance.

get_config()[source]

Returns the configuration of the layer.

Returns: Python dictionary with the layer’s parameters. Can be used to initialize and build another layer instance (of the same class), via from_config().

AttentiveFPConv

class molgraph.layers.AttentiveFPConv(molgraph.layers.GATConv)[source]

Node message passing step (“Atom embedding”) of AttentiveFP.

AttentiveFPConv inherits from GATConv and adds a GRU update. Also performs a initial linear transformation on node features if needed or desired.

For a complete implementation of AttentiveFP, add AttentiveFPReadout after AttentiveFPConv steps.

The implementation is based on Xiong et al. (2020) 20.

Example usage:

>>> 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 a model with AttentiveFPConv
>>> gnn_model = tf.keras.Sequential([
...     molgraph.layers.AttentiveFPConv(16, apply_initial_node_projection=True),
...     molgraph.layers.AttentiveFPConv(16)
... ])
>>> gnn_model(graph_tensor).node_feature.shape
TensorShape([5, 16])

Create a complete AttentiveFP model:

>>> 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],
... )
>>> attentive_fp = tf.keras.Sequential([
...     molgraph.layers.AttentiveFPConv(16, apply_initial_node_projection=True),
...     molgraph.layers.AttentiveFPConv(16),
...     molgraph.layers.AttentiveFPReadout(),
... ])
>>> attentive_fp(graph_tensor).shape
TensorShape([2, 16])

Pass the same GRU cell to each layer to perform weight sharing:

>>> 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],
... )
>>> gru_cell = tf.keras.layers.GRUCell(16) # same units as AttentiveFPConv
>>> gnn_model = tf.keras.Sequential([
...     molgraph.layers.AttentiveFPConv(
...         16, apply_initial_node_projection=True, gru_cell=gru_cell),
...     molgraph.layers.AttentiveFPConv(16, gru_cell=gru_cell)
... ])
>>> gnn_model(graph_tensor).node_feature.shape
TensorShape([5, 16])

Parameters

units (int, None) – Number of output units.
apply_initial_node_projection (bool) – Whether to perform an initial linear transformation on node features. Should be set to True for the first AttentiveFPConv layer in a sequence of AttentiveFPConv layers. Note that a node projection automatically occurs when units != node_feature_shape[-1]. Default to False.
gru_cell (tf.keras.layers.GRUCell, None) – For weight-tying of the update step (GRU step) provide a GRU cell. Default to None.
num_heads (int) – Number of attention heads. Default to 8.
merge_mode (str) – The strategy for merging the heads. Either of ‘concat’, ‘sum’, ‘mean’ or None. If set to None, ‘mean’ is used. Default to ‘concat’.
self_projection (bool) – Whether to apply self projection. Default to True.
normalization – (None, str, bool): Whether to apply layer normalization to the output. If batch normalization is desired, pass ‘batch_norm’. Default to None.
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.
attention_activation (tf.keras.activations.Activation, callable, str, None) – Activation function applied to the the attention scores. Default to ‘leaky_relu’.
activation (tf.keras.activations.Activation, callable, str, None) – Activation function applied to the output of the layer. 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. Default to tf.keras.initializers.TruncatedNormal(stddev=0.005).
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.
**kwargs –
Valid (optional) keyword arguments are:
- name (str): Name of the layer instance.
- update_step (tf.keras.layers.Layer): Applies post-processing step on the output (produced by _call). If passed, normalization, residual, activation and dropout parameters will be ignored. If None, a default post-processing step will be used (taking into consideration the aforementioned parameters). Default to None.

References

20: https://pubs.acs.org/doi/10.1021/acs.jmedchem.9b00959

build_from_signature(graph_tensor)[source]

Builds the layer based on a GraphTensor or a Spec.

Automatically invoked on first call().

Parameters: graph_tensor (GraphTensor, GraphTensor.Spec) – A graph tensor instance or a spec of a graph tensor instance.

call(graph_tensor)[source]

Wraps _call(), applying pre- and post-processing steps.

Automatically builds the layer on first call.

Parameters: graph_tensor (GraphTensor) – An graph tensor instance.
Returns: An updated graph tensor instance.
Return type: GraphTensor

compute_output_shape(input_shape)[source]

Computes the output shape of the layer based on an input shape.

Parameters: input_shape (tf.TensorShape) – The shape of a GraphTensor instance.
Returns: The shape corresponding to the outputted (updated) GraphTensor instance.
Return type: tf.TensorShape

compute_output_signature(input_signature)[source]

Computes the output signature of the layer based on an input signature.

Parameters: input_signature (GraphTensor.Spec) – The spec of a GraphTensor instance.
Returns: The spec corrsponding to the outputted (updated) GraphTensor instance.
Return type: GraphTensor.Spec

classmethod from_config(config)[source]

Initializes and builds the layer based on a configuration.

Parameters: config (dict) – A Python dictionary of parameters to initialize and build a layer instance. The config is usually obtained from get_config() of another layer instance (of the same class).
Returns: A layer instance.

get_config()[source]

Returns the configuration of the layer.

Returns: Python dictionary with the layer’s parameters. Can be used to initialize and build another layer instance (of the same class), via from_config().

Message-passing

MPNNConv

class molgraph.layers.MPNNConv(molgraph.layers.GNNLayer)[source]
Message passing neural network layer (MPNN)

Implementation is based on Gilmer et al. (2017) 21. In contrast to Gilmer et al. this implementation does not use weight tying by default; for neither the message function nor the update function. Furthermore, instead of the GRU-based update function, here, a simple fully-connected (dense) layer can be used too. Though default is GRU.

Example usage:
>>> 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],
... )
>>> gnn_model = tf.keras.Sequential([
...     molgraph.layers.MPNNConv(units=16),
...     molgraph.layers.MPNNConv(units=16),
...     molgraph.layers.MPNNConv(units=16),
...     molgraph.layers.Readout(),
... ])
>>> gnn_model(graph_tensor).shape
TensorShape([2, 16])
Pass the same GRU cell to each layer to perform weight sharing:
>>> 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],
... )
>>> gru_cell = tf.keras.layers.GRUCell(16)
>>> gnn_model = tf.keras.Sequential([
...     molgraph.layers.MPNNConv(units=16, update_fn=gru_cell),
...     molgraph.layers.MPNNConv(units=16, update_fn=gru_cell),
...     molgraph.layers.MPNNConv(units=16, update_fn=gru_cell),
...     molgraph.layers.Readout(),
... ])
>>> gnn_model(graph_tensor).shape
TensorShape([2, 16])
Parameters

units (int, None) – Number of output units.

update_mode (bool) – Specify what type of update will be performed. Either ‘dense’ or ‘gru’. Default to ‘gru’.

update_fn (tf.keras.layers.GRUCell, tf.keras.layers.Dense, None) – Optionally pass update function (GRUCell or Dense) for weight-tying of the update step (GRU step or Dense step). Default to None.

self_projection (bool) – Whether to apply self projection. Default to True.

normalization – (None, str, bool): Whether to apply layer normalization to the output. If batch normalization is desired, pass ‘batch_norm’. Default to None.

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.

update_activation (tf.keras.activations.Activation, callable, str, None) – Activation function used for the update function. Only relevant if ‘dense’ is passed to the update_mode argument. Default to None.

activation (tf.keras.activations.Activation, callable, str, None) – Activation function applied to the output of the layer. 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. Default to tf.keras.initializers.TruncatedNormal(stddev=0.005).

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.

**kwargs –
Valid (optional) keyword arguments are:

name (str): Name of the layer instance.

update_step (tf.keras.layers.Layer): Applies post-processing step on the output (produced by _call). If passed, normalization, residual, activation and dropout parameters will be ignored. If None, a default post-processing step will be used (taking into consideration the aforementioned parameters). Default to None.

References

21

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

call(graph_tensor)[source]

Wraps _call(), applying pre- and post-processing steps.

Automatically builds the layer on first call.

Parameters

graph_tensor (GraphTensor) – An graph tensor instance.

Returns

An updated graph tensor instance.

Return type

GraphTensor

compute_output_shape(input_shape)[source]

Computes the output shape of the layer based on an input shape.

Parameters

input_shape (tf.TensorShape) – The shape of a GraphTensor instance.

Returns

The shape corresponding to the outputted (updated) GraphTensor instance.

Return type

tf.TensorShape

classmethod from_config(config)[source]

Initializes and builds the layer based on a configuration.

Parameters

config (dict) – A Python dictionary of parameters to initialize and build a layer instance. The config is usually obtained from get_config() of another layer instance (of the same class).

Returns

A layer instance.

get_config()[source]

Returns the configuration of the layer.

Returns

Python dictionary with the layer’s parameters. Can be used to initialize and build another layer instance (of the same class), via from_config().

EdgeConv

class molgraph.layers.EdgeConv(tf.keras.layers.Layer)[source]
Edge convolutional layer, used to build DMPNN 22 and DGIN 23 like models.

Important:

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

Example usage:
>>> 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],
... )
>>> gnn_model = tf.keras.Sequential([
...     molgraph.layers.EdgeConv(16),                      # will produce 'edge_state' component
...     molgraph.layers.EdgeConv(16),                      # will produce 'edge_state' component
...     molgraph.layers.NodeReadout(target='edge_state'),  # target='edge_state' is default
...     molgraph.layers.Readout(),
... ])
>>> gnn_model(graph_tensor).shape
TensorShape([2, 16])
Parameters

units (int, None) – Number of output units.

update_mode (bool) – Specify what type of update will be performed. Either ‘dense’ or ‘gru’. Default to ‘gru’.

update_fn (tf.keras.layers.GRUCell, tf.keras.layers.Dense, None) – Optionally pass update function (GRUCell or Dense) for weight-tying of the update step (GRU step or Dense step). Default to None.

activation (tf.keras.activations.Activation, callable, str, None) – Activation function applied to the updated edge states. If None is set, either ‘relu’ (for update_mode=’dense’) or ‘tanh’ (for update_mode=’gru’) will be used. Default to None.

recurrent_activation (tf.keras.activations.Activation, callable, str, None) – Activation function applied to the recurrent step (only relevant if update_mode='gru'). Default to to ‘sigmoid’.

use_bias (bool) – Whether the layer should use biases. If None is set, either False or True will be used. Default to None.

kernel_initializer (tf.keras.initializers.Initializer, str) – Initializer function for the kernels. If None is set, either tf.keras.initializers.TruncatedNormal(stddev=0.005) (dense) or glorot_uniform (gru) 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.).

recurrent_initializer (tf.keras.initializers.Initializer, str) – Initializer function for the recurrent kernel (only relevant if update_mode='gru'). Default to ‘orthogonal’.

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.

recurrent_regularizer (tf.keras.regularizers.Regularizer, None) – Regularizer function applied to the recurrent kernel (only relevant if update_mode='gru'). 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.

recurrent_constraint (tf.keras.constraints.Constraint, None) – Constraint function applied to the recurrent kernel (only relevant if update_mode='gru'). Default to None.

References

22

https://arxiv.org/pdf/1904.01561.pdf

23

https://www.mdpi.com/1420-3049/26/20/6185

call(tensor)[source]

This is where the layer’s logic lives.

The call() method may not create state (except in its first invocation, wrapping the creation of variables or other resources in tf.init_scope()). It is recommended to create state, including tf.Variable instances and nested Layer instances,

in __init__(), or in the build() method that is

called automatically before call() executes for the first time.

Parameters

inputs –
Input tensor, or dict/list/tuple of input tensors. The first positional inputs argument is subject to special rules: - inputs must be explicitly passed. A layer cannot have zero

arguments, and inputs cannot be provided via the default value of a keyword argument.

NumPy array or Python scalar values in inputs get cast as tensors.

TF-Keras mask metadata is only collected from inputs.

Layers are built (build(input_shape) method) using shape info from inputs only.

input_spec compatibility is only checked against inputs.

Mixed precision input casting is only applied to inputs. If a layer has tensor arguments in *args or **kwargs, their casting behavior in mixed precision should be handled manually.

The SavedModel input specification is generated using inputs only.

Integration with various ecosystem packages like TFMOT, TFLite, TF.js, etc is only supported for inputs and not for tensors in positional and keyword arguments.

*args – Additional positional arguments. May contain tensors, although this is not recommended, for the reasons above.

**kwargs –
Additional keyword arguments. May contain tensors, although this is not recommended, for the reasons above. The following optional keyword arguments are reserved: - training: Boolean scalar tensor of Python boolean indicating

whether the call is meant for training or inference.

mask: Boolean input mask. If the layer’s call() method takes a mask argument, its default value will be set to the mask generated for inputs by the previous layer (if input did come from a layer that generated a corresponding mask, i.e. if it came from a TF-Keras layer with masking support).

Returns

A tensor or list/tuple of tensors.

compute_output_shape(input_shape)[source]

Computes the output shape of the layer.

This method will cause the layer’s state to be built, if that has not happened before. This requires that the layer will later be used with inputs that match the input shape provided here.

Parameters

input_shape – Shape tuple (tuple of integers) or tf.TensorShape, or structure of shape tuples / tf.TensorShape instances (one per output tensor of the layer). Shape tuples can include None for free dimensions, instead of an integer.

Returns

A tf.TensorShape instance or structure of tf.TensorShape instances.

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.

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.

Geometric

DTNNConv

class molgraph.layers.DTNNConv(molgraph.layers.GNNLayer)[source]

Deep Tensor Neural Network (DTNN).

Implementation is based on Schütt et al. (2017a) 24.

Example usage:

>>> graph_tensor = molgraph.GraphTensor(
...     sizes=[2, 3],
...     node_feature=[[1., 0.], [1., 0.], [1., 0.], [1., 0.], [0., 1.]],
...     edge_src=[1, 0, 3, 4, 2, 4, 3, 2],
...     edge_dst=[0, 1, 2, 2, 3, 3, 4, 4],
...     # edge_feature encodes e.g. distances between edge_src and edge_dst
...     edge_feature=[0.3, 0.3, 0.1, 0.2, 0.1, 0.4, 0.4, 0.2],
... )
>>> gnn_model = tf.keras.Sequential([
...     molgraph.layers.DTNNConv(units=16),
...     molgraph.layers.DTNNConv(units=16),
...     molgraph.layers.DTNNConv(units=16),
...     molgraph.layers.Readout(),
... ])
>>> gnn_model(graph_tensor).shape
TensorShape([2, 16])

Parameters

units (int, None) – The number of output units.
distance_min (float) – The smallest center (mean) to be used for the radial basis function. I.e., it defines the minimum distance between atom pairs; or the minimum electrostatic interaction between nuclei, in the case of Coulomb values. Default to -1.0.
distance_max (float) – The largest center (mean) to be used for the radial basis function. I.e., it defines the maximum distance between atom pairs; or the maximum electrostatic interaction between nuclei, in the case of Coulomb values. Default to 18.0.
distance_granularity (float) – The distance between each center (mean) of the radial basis function. The smaller the granularity, the more centers will be used. Default to 0.1.
rbf_stddev (float, str) – The standard deviation of the radial basis function. If ‘auto’, ‘distance_granularity’ will be used as standard deviation. Default to ‘auto’.
self_projection (bool) – Whether to apply self projection. Default to True.
normalization – (None, str, bool): Whether to apply layer normalization to the output. If batch normalization is desired, pass ‘batch_norm’. Default to None.
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 output of the layer. Default to None.
use_bias (bool) – Whether the layer should use biases. Default to True.
kernel_initializer (tf.keras.initializers.Initializer, str) – Initializer function for the kernels. Default to tf.keras.initializers.GlorotUniform().
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.
**kwargs –
Valid (optional) keyword arguments are:
- name (str): Name of the layer instance.
- update_step (tf.keras.layers.Layer): Applies post-processing step on the output (produced by _call). If passed, normalization, residual, activation and dropout parameters will be ignored. If None, a default post-processing step will be used (taking into consideration the aforementioned parameters). Default to None.

References

24: https://www.nature.com/articles/ncomms13890

call(graph_tensor)[source]

Wraps _call(), applying pre- and post-processing steps.

Automatically builds the layer on first call.

Parameters: graph_tensor (GraphTensor) – An graph tensor instance.
Returns: An updated graph tensor instance.
Return type: GraphTensor

compute_output_shape(input_shape)[source]

Computes the output shape of the layer based on an input shape.

Parameters: input_shape (tf.TensorShape) – The shape of a GraphTensor instance.
Returns: The shape corresponding to the outputted (updated) GraphTensor instance.
Return type: tf.TensorShape

classmethod from_config(config)[source]

Initializes and builds the layer based on a configuration.

Parameters: config (dict) – A Python dictionary of parameters to initialize and build a layer instance. The config is usually obtained from get_config() of another layer instance (of the same class).
Returns: A layer instance.

get_config()[source]

Returns the configuration of the layer.

Returns: Python dictionary with the layer’s parameters. Can be used to initialize and build another layer instance (of the same class), via from_config().

GCFConv

class molgraph.layers.GCFConv(molgraph.layers.GNNLayer)[source]

(Graph) continuous filter convolution layer ((G)CFConv).

Implementation is based on Schütt et al. (2017b) 25.

Operates on 3D molecular graphs (encoding distance geometry).

Example usage:

>>> graph_tensor = molgraph.GraphTensor(
...     sizes=[2, 3],
...     node_feature=[[1., 0.], [1., 0.], [1., 0.], [1., 0.], [0., 1.]],
...     edge_src=[1, 0, 3, 4, 2, 4, 3, 2],
...     edge_dst=[0, 1, 2, 2, 3, 3, 4, 4],
...     # edge_feature encodes e.g. distances between edge_src and edge_dst
...     edge_feature=[0.3, 0.3, 0.1, 0.2, 0.1, 0.4, 0.4, 0.2],
... )
>>> gnn_model = tf.keras.Sequential([
...     molgraph.layers.GCFConv(units=16),
...     molgraph.layers.GCFConv(units=16),
...     molgraph.layers.GCFConv(units=16),
...     molgraph.layers.Readout(),
... ])
>>> gnn_model(graph_tensor).shape
TensorShape([2, 16])

Parameters

units (int, None) – The number of output units.
distance_min (float) – The smallest center (mean) to be used for the radial basis function. I.e., it defines the minimum distance between atom pairs; or the minimum electrostatic interaction between nuclei, in the case of Coulomb values. Default to -1.0.
distance_max (float) – The largest center (mean) to be used for the radial basis function. I.e., it defines the maximum distance between atom pairs; or the maximum electrostatic interaction between nuclei, in the case of Coulomb values. Default to 18.0.
distance_granularity (float) – The distance between each center (mean) of the radial basis function. The smaller the granularity, the more centers will be used. Default to 0.1.
rbf_stddev (float, str) – The standard deviation of the radial basis function. If ‘auto’, ‘distance_granularity’ will be used as standard deviation. Default to ‘auto’.
self_projection (bool) – Whether to apply self projection. Default to True.
normalization – (None, str, bool): Whether to apply layer normalization to the output. If batch normalization is desired, pass ‘batch_norm’. Default to None.
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 output of the layer. Default to None.
use_bias (bool) – Whether the layer should use biases. Default to True.
kernel_initializer (tf.keras.initializers.Initializer, str) – Initializer function for the kernels. Default to tf.keras.initializers.GlorotUniform().
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.
**kwargs –
Valid (optional) keyword arguments are:
- name (str): Name of the layer instance.
- update_step (tf.keras.layers.Layer): Applies post-processing step on the output (produced by _call). If passed, normalization, residual, activation and dropout parameters will be ignored. If None, a default post-processing step will be used (taking into consideration the aforementioned parameters). Default to None.

References

25: https://arxiv.org/pdf/1706.08566.pdf

call(graph_tensor)[source]

Wraps _call(), applying pre- and post-processing steps.

Automatically builds the layer on first call.

Parameters: graph_tensor (GraphTensor) – An graph tensor instance.
Returns: An updated graph tensor instance.
Return type: GraphTensor

compute_output_shape(input_shape)[source]

Computes the output shape of the layer based on an input shape.

Parameters: input_shape (tf.TensorShape) – The shape of a GraphTensor instance.
Returns: The shape corresponding to the outputted (updated) GraphTensor instance.
Return type: tf.TensorShape

classmethod from_config(config)[source]

Initializes and builds the layer based on a configuration.

Parameters: config (dict) – A Python dictionary of parameters to initialize and build a layer instance. The config is usually obtained from get_config() of another layer instance (of the same class).
Returns: A layer instance.

get_config()[source]

Returns the configuration of the layer.

Returns: Python dictionary with the layer’s parameters. Can be used to initialize and build another layer instance (of the same class), via from_config().

Readout

SegmentPoolingReadout

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

Segmentation pooling for graph readout.

Alias: Readout

Example usage:

>>> graph_tensor = molgraph.GraphTensor(
...     sizes=[2, 3],
...     node_feature=[[1., 0.], [1., 0.], [1., 0.], [1., 0.], [0., 1.]],
...     edge_src=[1, 0, 3, 4, 2, 4, 3, 2],
...     edge_dst=[0, 1, 2, 2, 3, 3, 4, 4],
... )
>>> model = tf.keras.Sequential([
...     # molgraph.layers.GCNConv(4),
...     molgraph.layers.Readout('mean') # alias for SegmentPoolingReadout
... ])
>>> model(graph_tensor)
<tf.Tensor: shape=(2, 2), dtype=float32, numpy=
array([[1.        , 0.        ],
       [0.6666667 , 0.33333334]], dtype=float32)>

Parameters: mode (str) – What type of pooling should be performed. Either of ‘mean’, ‘max’ or ‘sum’. Specifically, performs a tf.math.segment_mean, tf.math.segment_max or tf.math.segment_sum (respectively) on the node_feature field of the inputted graph tensor. Defaults to ‘mean’.

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 tf.Tensor or tf.RaggedTensor based on the node_feature field of the inputted GraphTensor.

compute_output_shape(input_shape)[source]

Computes the output shape of the layer.

This method will cause the layer’s state to be built, if that has not happened before. This requires that the layer will later be used with inputs that match the input shape provided here.

Parameters: input_shape – Shape tuple (tuple of integers) or tf.TensorShape, or structure of shape tuples / tf.TensorShape instances (one per output tensor of the layer). Shape tuples can include None for free dimensions, instead of an integer.
Returns: A tf.TensorShape instance or structure of tf.TensorShape instances.

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.

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.

TransformerEncoderReadout

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

Transformer encoder layer for graph readout.

Example usage:

>>> graph_tensor = molgraph.GraphTensor(
...     sizes=[2, 3],
...     node_feature=[[1., 0.], [1., 0.], [1., 0.], [1., 0.], [0., 1.]],
...     edge_src=[1, 0, 3, 4, 2, 4, 3, 2],
...     edge_dst=[0, 1, 2, 2, 3, 3, 4, 4],
... )
>>> model = tf.keras.Sequential([
...     # molgraph.layers.GCNConv(4),
...     molgraph.layers.TransformerEncoderReadout()
... ])
>>> model(graph_tensor).shape
TensorShape([2, 2])

Parameters

hidden_units (int) – Number of hidden units (of the feedforward network).
num_heads (int) – Number of attention heads. Default to 8.
activation (str, tf.keras.activations.Activation, None) – The activation function applied to the output. Default to ‘relu’.

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 tf.Tensor or tf.RaggedTensor based on the node_feature field of the inputted GraphTensor.

compute_output_shape(input_shape)[source]

Computes the output shape of the layer.

This method will cause the layer’s state to be built, if that has not happened before. This requires that the layer will later be used with inputs that match the input shape provided here.

Parameters: input_shape – Shape tuple (tuple of integers) or tf.TensorShape, or structure of shape tuples / tf.TensorShape instances (one per output tensor of the layer). Shape tuples can include None for free dimensions, instead of an integer.
Returns: A tf.TensorShape instance or structure of tf.TensorShape instances.

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.

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.

SetGatherReadout

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

Set-to-set layer for graph readout.

Implementation based on Gilmer et al. (2017) 26 and Vinyals et al. (2016) 27.

Example usage:

>>> graph_tensor = molgraph.GraphTensor(
...     sizes=[2, 3],
...     node_feature=[[1., 0.], [1., 0.], [1., 0.], [1., 0.], [0., 1.]],
...     edge_src=[1, 0, 3, 4, 2, 4, 3, 2],
...     edge_dst=[0, 1, 2, 2, 3, 3, 4, 4],
... )
>>> model = tf.keras.Sequential([
...     # molgraph.layers.GCNConv(2),
...     molgraph.layers.SetGatherReadout()
... ])
>>> # Note: SetGatherReadout doubles output dim (from 2 to 4)
>>> model(graph_tensor).shape
TensorShape([2, 4])

Parameters: steps (int) – Number of LSTM steps. Default to 8.

References

26: https://arxiv.org/pdf/1704.01212.pdf
27: https://arxiv.org/pdf/1511.06391.pdf

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 tf.Tensor or tf.RaggedTensor based on the node_feature field of the inputted GraphTensor.

compute_output_shape(input_shape)[source]

Computes the output shape of the layer.

This method will cause the layer’s state to be built, if that has not happened before. This requires that the layer will later be used with inputs that match the input shape provided here.

Parameters: input_shape – Shape tuple (tuple of integers) or tf.TensorShape, or structure of shape tuples / tf.TensorShape instances (one per output tensor of the layer). Shape tuples can include None for free dimensions, instead of an integer.
Returns: A tf.TensorShape instance or structure of tf.TensorShape instances.

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.

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.

NodeReadout

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

Aggregates edge states to associated nodes.

Example usage:

>>> 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],
... )
>>> gnn_model = tf.keras.Sequential([
...     molgraph.layers.EdgeConv(16),                           # will produce 'edge_state' component
...     molgraph.layers.EdgeConv(16),                           # will produce 'edge_state' component
...     molgraph.layers.NodeReadout(target='edge_state'),       # target='edge_state' is default
... ])
>>> gnn_model(graph_tensor).node_feature.shape
TensorShape([5, 16])

Parameters

target (str) – Specifies which component to aggregate. Default to ‘edge_state’ which is the component produced by molgraph.layers.EdgeConv.
apply_transform (bool) – Whether to perform a transformaton after the aggregation. Default to False.
dense_kwargs (None, dict) – Parameters to be passed to the dense layer in the transformation. Only relevant if apply_transform=True. If None is passed, units is set to input_shape[-1] and activation is set to relu. Default to None.

call(tensor)[source]

This is where the layer’s logic lives.

The call() method may not create state (except in its first invocation, wrapping the creation of variables or other resources in tf.init_scope()). It is recommended to create state, including tf.Variable instances and nested Layer instances,

in __init__(), or in the build() method that is

called automatically before call() executes for the first time.

Parameters

inputs –
Input tensor, or dict/list/tuple of input tensors. The first positional inputs argument is subject to special rules: - inputs must be explicitly passed. A layer cannot have zero

arguments, and inputs cannot be provided via the default value of a keyword argument.
- NumPy array or Python scalar values in inputs get cast as tensors.
- TF-Keras mask metadata is only collected from inputs.
- Layers are built (build(input_shape) method) using shape info from inputs only.
- input_spec compatibility is only checked against inputs.
- Mixed precision input casting is only applied to inputs. If a layer has tensor arguments in *args or **kwargs, their casting behavior in mixed precision should be handled manually.
- The SavedModel input specification is generated using inputs only.
- Integration with various ecosystem packages like TFMOT, TFLite, TF.js, etc is only supported for inputs and not for tensors in positional and keyword arguments.
*args – Additional positional arguments. May contain tensors, although this is not recommended, for the reasons above.
**kwargs –
Additional keyword arguments. May contain tensors, although this is not recommended, for the reasons above. The following optional keyword arguments are reserved: - training: Boolean scalar tensor of Python boolean indicating

whether the call is meant for training or inference.
- mask: Boolean input mask. If the layer’s call() method takes a mask argument, its default value will be set to the mask generated for inputs by the previous layer (if input did come from a layer that generated a corresponding mask, i.e. if it came from a TF-Keras layer with masking support).

Returns

A tensor or list/tuple of tensors.

compute_output_shape(input_shape)[source]

Computes the output shape of the layer.

This method will cause the layer’s state to be built, if that has not happened before. This requires that the layer will later be used with inputs that match the input shape provided here.

Parameters: input_shape – Shape tuple (tuple of integers) or tf.TensorShape, or structure of shape tuples / tf.TensorShape instances (one per output tensor of the layer). Shape tuples can include None for free dimensions, instead of an integer.
Returns: A tf.TensorShape instance or structure of tf.TensorShape instances.

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.

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.

AttentiveFPReadout

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

Readout step (“Molecule embedding”) of AttentiveFP.

For a complete implementation of AttentiveFP, add a number of AttentiveFPConv steps before AttentiveFPReadout (see below).

The implementation is based on Xiong et al. (2020) 28.

Example usage:

Create a complete AttentiveFP model:

>>> 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 a model with AttentiveFPConv
>>> attentive_fp = tf.keras.Sequential([
...     molgraph.layers.AttentiveFPConv(16, apply_initial_node_projection=True),
...     molgraph.layers.AttentiveFPConv(16),
...     # ...
...     molgraph.layers.AttentiveFPReadout(steps=4),
... ])
>>> attentive_fp(graph_tensor).shape
TensorShape([2, 16])

Parameters

steps (int) – Number of aggregation steps to perform. Default to 4.
message_step (molgraph.layers.attentional.gat_conv.GATConv, None) – The message passing step.
update_step (tf.keras.layers.GRUCell, None) – The update step.
final_node_projection (tf.keras.layers.Dense, None) – The final projection applied to the output.

References

28: https://pubs.acs.org/doi/10.1021/acs.jmedchem.9b00959

call(tensor)[source]

This is where the layer’s logic lives.

The call() method may not create state (except in its first invocation, wrapping the creation of variables or other resources in tf.init_scope()). It is recommended to create state, including tf.Variable instances and nested Layer instances,

in __init__(), or in the build() method that is

called automatically before call() executes for the first time.

Parameters

inputs –
Input tensor, or dict/list/tuple of input tensors. The first positional inputs argument is subject to special rules: - inputs must be explicitly passed. A layer cannot have zero

arguments, and inputs cannot be provided via the default value of a keyword argument.
- NumPy array or Python scalar values in inputs get cast as tensors.
- TF-Keras mask metadata is only collected from inputs.
- Layers are built (build(input_shape) method) using shape info from inputs only.
- input_spec compatibility is only checked against inputs.
- Mixed precision input casting is only applied to inputs. If a layer has tensor arguments in *args or **kwargs, their casting behavior in mixed precision should be handled manually.
- The SavedModel input specification is generated using inputs only.
- Integration with various ecosystem packages like TFMOT, TFLite, TF.js, etc is only supported for inputs and not for tensors in positional and keyword arguments.
*args – Additional positional arguments. May contain tensors, although this is not recommended, for the reasons above.
**kwargs –
Additional keyword arguments. May contain tensors, although this is not recommended, for the reasons above. The following optional keyword arguments are reserved: - training: Boolean scalar tensor of Python boolean indicating

whether the call is meant for training or inference.
- mask: Boolean input mask. If the layer’s call() method takes a mask argument, its default value will be set to the mask generated for inputs by the previous layer (if input did come from a layer that generated a corresponding mask, i.e. if it came from a TF-Keras layer with masking support).

Returns

A tensor or list/tuple of tensors.

compute_output_shape(input_shape)[source]

Computes the output shape of the layer.

This method will cause the layer’s state to be built, if that has not happened before. This requires that the layer will later be used with inputs that match the input shape provided here.

Parameters: input_shape – Shape tuple (tuple of integers) or tf.TensorShape, or structure of shape tuples / tf.TensorShape instances (one per output tensor of the layer). Shape tuples can include None for free dimensions, instead of an integer.
Returns: A tf.TensorShape instance or structure of tf.TensorShape instances.

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.

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.

Positional encoding

LaplacianPositionalEncodig

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

Laplacian positional encoding.

Implementation based on Dwivedi et al. (2021) 29 and Belkin et al. (2003) 30.

Example usage:

>>> graph_tensor = molgraph.GraphTensor(
...     sizes=[2, 3],
...     node_feature=[[1., 0.], [1., 0.], [1., 0.], [1., 0.], [0., 1.]],
...     edge_src=[1, 0, 3, 4, 2, 4, 3, 2],
...     edge_dst=[0, 1, 2, 2, 3, 3, 4, 4],
... )
>>> model = tf.keras.Sequential([
...     molgraph.layers.LaplacianPositionalEncoding(16),
... ])
>>> graph_tensor = model(graph_tensor)
>>> graph_tensor.node_feature != 1.0
<tf.Tensor: shape=(5, 2), dtype=bool, numpy=
array([[ True,  True],
       [ True,  True],
       [ True,  True],
       [ True,  True],
       [ True,  True]])>

Parameters

dim (int) – The dimension of the positional encoding. Default to 8.
activation (tf.keras.activations.Activation, callable, str, None) – Activation function applied to the output of the layer. Default to None.
use_bias (bool) – Whether the layer should use biases. Default to False.
kernel_initializer (tf.keras.initializers.Initializer, str) – Initializer function for the kernel. Default to tf.keras.initializers.TruncatedNormal(stddev=0.005).
bias_initializer (tf.keras.initializers.Initializer, str) – Initializer function for the bias. Default to tf.keras.initializers.Constant(0.).
kernel_regularizer (tf.keras.regularizers.Regularizer, None) – Regularizer function applied to the kernel. Default to None.
bias_regularizer (tf.keras.regularizers.Regularizer, None) – Regularizer function applied to the bias. 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 kernel. Default to None.
bias_constraint (tf.keras.constraints.Constraint, None) – Constraint function applied to the bias. Default to None.

References

29: https://arxiv.org/pdf/2012.09699.pdf
30: https://ieeexplore.ieee.org/document/6789755

call(tensor, training=None)[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) – A graph tensor which serves as input to the layer.
Returns: A graph tensor with updated node features.
Return type: GraphTensor

compute_output_shape(input_shape)[source]

Computes the output shape of the layer.

This method will cause the layer’s state to be built, if that has not happened before. This requires that the layer will later be used with inputs that match the input shape provided here.

Parameters: input_shape – Shape tuple (tuple of integers) or tf.TensorShape, or structure of shape tuples / tf.TensorShape instances (one per output tensor of the layer). Shape tuples can include None for free dimensions, instead of an integer.
Returns: A tf.TensorShape instance or structure of tf.TensorShape instances.

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.

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.

Preprocessing

MinMaxScaling

class molgraph.layers.MinMaxScaling(layers.experimental.preprocessing.PreprocessingLayer)[source]

Min-max scaling between a specified range.

Specify, as keyword argument only, MinMaxScaling(feature='node_feature') to perform min-max scaling on the node_feature component of the GraphTensor, or, MinMaxScaling(feature='edge_feature') to perform min-max scaling on the edge_feature component of the GraphTensor. If not specified, the node_feature component will be considered.

Example usage:

>>> graph_tensor = molgraph.GraphTensor(
...     sizes=[2, 3],
...     node_feature=[[2., .5], [2., 0.], [2., 0.], [2., .5], [0., 2.]],
...     edge_src=[1, 0, 3, 4, 2, 4, 3, 2],
...     edge_dst=[0, 1, 2, 2, 3, 3, 4, 4],
... )
>>> preprocessing = molgraph.layers.MinMaxScaling(
...    feature='node_feature')
>>> preprocessing.adapt(graph_tensor)
>>> model = tf.keras.Sequential([preprocessing])
>>> graph_tensor = model(graph_tensor)
>>> graph_tensor.node_feature
<tf.Tensor: shape=(5, 2), dtype=float32, numpy=
array([[1.  , 0.25],
       [1.  , 0.  ],
       [1.  , 0.  ],
       [1.  , 0.25],
       [0.  , 1.  ]], dtype=float32)>

Adapt layer on a tf.data.Dataset:

>>> graph_tensor = molgraph.GraphTensor(
...     sizes=[2, 3],
...     node_feature=[[2., .5], [2., 0.], [2., 0.], [2., .5], [0., 2.]],
...     edge_src=[1, 0, 3, 4, 2, 4, 3, 2],
...     edge_dst=[0, 1, 2, 2, 3, 3, 4, 4],
... )
>>> ds = tf.data.Dataset.from_tensor_slices(graph_tensor).batch(2)
>>> preprocessing = molgraph.layers.MinMaxScaling(
...    feature='node_feature')
>>> preprocessing.adapt(ds)
>>> model = tf.keras.Sequential([preprocessing])
>>> output = model.predict(ds, verbose=0)
>>> output.node_feature
<tf.Tensor: shape=(5, 2), dtype=float32, numpy=
array([[1.  , 0.25],
       [1.  , 0.  ],
       [1.  , 0.  ],
       [1.  , 0.25],
       [0.  , 1.  ]], dtype=float32)>

Parameters

feature_range (tuple) – The range of values of transformed data.
minimum (tf.Tensor, None) – The minimum values of the original data; used to transform the data into the given feature_range. If None, the layer has to be adapted via e.g. adapt(). Default to None.
maximum (tf.Tensor, None) – The maximum values of the original data; used to transform the data into the given feature_range. If None, the layer has to be adapted via e.g. adapt(). Default to None.
**kwargs – Specify the relevant feature. Default to node_feature. The reminaing kwargs are passed to the parent class.

adapt(data, batch_size=None, steps=None)[source]

Adapts the layer to data.

When adapting the layer to the data, build() will be called automatically (to initialize the relevant attributes). After adaption, the layer is finalized and ready to be used.

Parameters

data (GraphTensor, tf.data.Dataset) – Data to be used to adapt the layer. Can be either a GraphTensor directly or a tf.data.Dataset constructed from a GraphTensor.
batch_size (int, None) – The batch size to be used during adaption. Default to None.
steps (int, None) – The number of steps of adaption. If None, the number of samples divided by the batch_size is used. Default to None.

call(data)[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: data (GraphTensor) – Input to the layer.
Returns: A GraphTensor with updated features. Either the node_feature component or the edge_feature component (of the GraphTensor) are updated.
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.

StandardScaling

class molgraph.layers.StandardScaling(layers.experimental.preprocessing.PreprocessingLayer)[source]

Standard scaling, via centering and standardization.

Specify, as keyword argument only, StandardScaling(feature='node_feature') to perform standard scaling on the node_feature component of the GraphTensor, or, StandardScaling(feature='edge_feature') to perform standard scaling on the edge_feature component of the GraphTensor. If not specified, the node_feature component will be considered.

Examples:

Adapt layer on GraphTensor directly:

>>> graph_tensor = molgraph.GraphTensor(
...     sizes=[2, 3],
...     node_feature=[[2., .5], [2., 0.], [2., 0.], [2., .5], [0., 2.]],
...     edge_src=[1, 0, 3, 4, 2, 4, 3, 2],
...     edge_dst=[0, 1, 2, 2, 3, 3, 4, 4],
... )
>>> preprocessing = molgraph.layers.StandardScaling(
...    feature='node_feature')
>>> preprocessing.adapt(graph_tensor)
>>> model = tf.keras.Sequential([preprocessing,])
>>> graph_tensor = model(graph_tensor)
>>> graph_tensor.node_feature
<tf.Tensor: shape=(5, 2), dtype=float32, numpy=
array([[ 0.49999997, -0.1360828 ],
       [ 0.49999997, -0.8164967 ],
       [ 0.49999997, -0.8164967 ],
       [ 0.49999997, -0.1360828 ],
       [-2.        ,  1.9051588 ]], dtype=float32)>

Adapt layer on a tf.data.Dataset:

>>> graph_tensor = molgraph.GraphTensor(
...     sizes=[2, 3],
...     node_feature=[[2., .5], [2., 0.], [2., 0.], [2., .5], [0., 2.]],
...     edge_src=[1, 0, 3, 4, 2, 4, 3, 2],
...     edge_dst=[0, 1, 2, 2, 3, 3, 4, 4],
... )
>>> ds = tf.data.Dataset.from_tensor_slices(graph_tensor).batch(2)
>>> preprocessing = molgraph.layers.StandardScaling(
...    feature='node_feature')
>>> preprocessing.adapt(ds)
>>> model = tf.keras.Sequential([preprocessing])
>>> output = model.predict(ds, verbose=0)
>>> output.node_feature
<tf.Tensor: shape=(5, 2), dtype=float32, numpy=
array([[ 0.49999997, -0.1360828 ],
       [ 0.49999997, -0.8164967 ],
       [ 0.49999997, -0.8164967 ],
       [ 0.49999997, -0.1360828 ],
       [-2.        ,  1.9051588 ]], dtype=float32)>

Parameters

mean (tf.Tensor, None) – The means of the original data; used to transform the data. If None, the layer has to be adapted via e.g. adapt(). Default to None.
variance (tf.Tensor, None) – The variances of the original data; used to transform the data. If None, the layer has to be adapted via e.g. adapt(). Default to None.
variance_threshold (int, float, None) – The threshold for removing features, based on the variance of the original data. If None, variance thresholding will not be performed. Default to None.
**kwargs – Specify the relevant feature. Default to node_feature. The reminaing kwargs are passed to the parent class.

adapt(data, batch_size=None, steps=None)[source]

Adapts the layer to data.

When adapting the layer to the data, build() will be called automatically (to initialize the relevant attributes). After adaption, the layer is finalized and ready to be used.

Parameters

data (GraphTensor, tf.data.Dataset) – Data to be used to adapt the layer. Can be either a GraphTensor directly or a tf.data.Dataset constructed from a GraphTensor.
batch_size (int, None) – The batch size to be used during adaption. Default to None.
steps (int, None) – The number of steps of adaption. If None, the number of samples divided by the batch_size is used. Default to None.

call(data)[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: data (GraphTensor) – Input to the layer.
Returns: A GraphTensor with updated features. Either the node_feature component or the edge_feature component (of the GraphTensor) are updated.
Return type: GraphTensor

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.

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.

VarianceThreshold

class molgraph.layers.VarianceThreshold(molgraph.layers.StandardScaling)[source]

Variance thresholding.

Uses the StandardScaling layer but ignores the normalization when calling the layer.

Specify, as keyword argument only, VarianceThreshold(feature='node_feature') to perform variance thresholding on the node_feature component of the GraphTensor, or, VarianceThreshold(feature='edge_feature') to perform variance thresholding on the edge_feature component of the GraphTensor. If not specified, the node_feature component will be considered.

Example usage:

>>> graph_tensor = molgraph.GraphTensor(
...     sizes=[2, 3],
...     node_feature=[[2., .5], [2., 0.], [2., 0.], [2., .5], [0., 2.]],
...     edge_src=[1, 0, 3, 4, 2, 4, 3, 2],
...     edge_dst=[0, 1, 2, 2, 3, 3, 4, 4],
... )
>>> preprocessing = molgraph.layers.VarianceThreshold(
...    feature='node_feature', variance_threshold=0.6)
>>> preprocessing.adapt(graph_tensor)
>>> model = tf.keras.Sequential([preprocessing])
>>> graph_tensor = model(graph_tensor)
>>> graph_tensor.node_feature
<tf.Tensor: shape=(5, 1), dtype=float32, numpy=
array([[2.],
       [2.],
       [2.],
       [2.],
       [0.]], dtype=float32)>

Adapt layer on a tf.data.Dataset:

>>> graph_tensor = molgraph.GraphTensor(
...     sizes=[2, 3],
...     node_feature=[[2., .5], [2., 0.], [2., 0.], [2., .5], [0., 2.]],
...     edge_src=[1, 0, 3, 4, 2, 4, 3, 2],
...     edge_dst=[0, 1, 2, 2, 3, 3, 4, 4],
... )
>>> ds = tf.data.Dataset.from_tensor_slices(graph_tensor).batch(2)
>>> preprocessing = molgraph.layers.VarianceThreshold(
...    feature='node_feature', variance_threshold=0.6)
>>> preprocessing.adapt(ds)
>>> model = tf.keras.Sequential([preprocessing])
>>> output = model.predict(ds, verbose=0)
>>> output.node_feature
<tf.Tensor: shape=(5, 1), dtype=float32, numpy=
array([[2.],
       [2.],
       [2.],
       [2.],
       [0.]], dtype=float32)>

Parameters

variance_threshold (int, float, None) – The threshold for removing features, based on the variance of the features over the graph. If None, variance thresholding will not be performed. Default to 0.001.
mean (tf.Tensor, None) – The mean of the features. Default to None.
variance (tf.Tensor, None) – The variance of the features. Default to None.
**kwargs – Specify the relevant feature. Default to node_feature. The reminaing kwargs are passed to the parent class.

call(data)[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: data (GraphTensor) – Input to the layer.
Returns: A GraphTensor with updated features. Either the node_feature component or the edge_feature component (of the GraphTensor) are updated.
Return type: GraphTensor

adapt(data, batch_size=None, steps=None)[source]

Adapts the layer to data.

When adapting the layer to the data, build() will be called automatically (to initialize the relevant attributes). After adaption, the layer is finalized and ready to be used.

Parameters

data (GraphTensor, tf.data.Dataset) – Data to be used to adapt the layer. Can be either a GraphTensor directly or a tf.data.Dataset constructed from a GraphTensor.
batch_size (int, None) – The batch size to be used during adaption. Default to None.
steps (int, None) – The number of steps of adaption. If None, the number of samples divided by the batch_size is used. Default to None.

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.

CenterScaling

class molgraph.layers.CenterScaling(layers.experimental.preprocessing.PreprocessingLayer)[source]

Centering.

Specify, as keyword argument only, CenterScaling(feature='node_feature') to perform center scaling on the node_feature component of the GraphTensor, or, CenterScaling(feature='edge_feature') to perform center scaling on the edge_feature component of the GraphTensor. If not specified, the node_feature component will be considered.

Example usage:

>>> graph_tensor = molgraph.GraphTensor(
...     sizes=[2, 3],
...     node_feature=[[1., 0.], [1., 0.], [1., 0.], [1., 0.], [0., 1.]],
...     edge_src=[1, 0, 3, 4, 2, 4, 3, 2],
...     edge_dst=[0, 1, 2, 2, 3, 3, 4, 4],
... )
>>> preprocessing = molgraph.layers.CenterScaling(
...    feature='node_feature')
>>> preprocessing.adapt(graph_tensor)
>>> model = tf.keras.Sequential([preprocessing])
>>> graph_tensor = model(graph_tensor)
>>> graph_tensor.node_feature
<tf.Tensor: shape=(5, 2), dtype=float32, numpy=
array([[ 0.19999999, -0.2       ],
       [ 0.19999999, -0.2       ],
       [ 0.19999999, -0.2       ],
       [ 0.19999999, -0.2       ],
       [-0.8       ,  0.8       ]], dtype=float32)>

Adapt layer on a tf.data.Dataset:

>>> graph_tensor = molgraph.GraphTensor(
...     sizes=[2, 3],
...     node_feature=[[1., 0.], [1., 0.], [1., 0.], [1., 0.], [0., 1.]],
...     edge_src=[1, 0, 3, 4, 2, 4, 3, 2],
...     edge_dst=[0, 1, 2, 2, 3, 3, 4, 4],
... )
>>> ds = tf.data.Dataset.from_tensor_slices(graph_tensor).batch(2)
>>> preprocessing = molgraph.layers.CenterScaling(
...    feature='node_feature')
>>> preprocessing.adapt(ds)
>>> model = tf.keras.Sequential([preprocessing])
>>> output = model.predict(ds, verbose=0)
>>> output.node_feature
<tf.Tensor: shape=(5, 2), dtype=float32, numpy=
array([[ 0.19999999, -0.2       ],
       [ 0.19999999, -0.2       ],
       [ 0.19999999, -0.2       ],
       [ 0.19999999, -0.2       ],
       [-0.8       ,  0.8       ]], dtype=float32)>

Parameters

mean (tf.Tensor, None) – The means of the original data; used to transform the data. If None, the layer has to be adapted via e.g. adapt(). Default to None.
**kwargs – Specify the relevant feature. Default to node_feature. The reminaing kwargs are passed to the parent class.

adapt(data, batch_size=None, steps=None)[source]

Adapts the layer to data.

When adapting the layer to the data, build() will be called automatically (to initialize the relevant attributes). After adaption, the layer is finalized and ready to be used.

Parameters

data (GraphTensor, tf.data.Dataset) – Data to be used to adapt the layer. Can be either a GraphTensor directly or a tf.data.Dataset constructed from a GraphTensor.
batch_size (int, None) – The batch size to be used during adaption. Default to None.
steps (int, None) – The number of steps of adaption. If None, the number of samples divided by the batch_size is used. Default to None.

call(data)[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: data (GraphTensor) – Input to the layer.
Returns: A GraphTensor with updated features. Either the node_feature component or the edge_feature component (of the GraphTensor) are updated.
Return type: GraphTensor

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.

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.

EmbeddingLookup

class molgraph.layers.EmbeddingLookup(tensorflow.keras.layers.StringLookup)[source]

A loookup layer and embedding layer in combination.

Specify, as keyword argument only, EmbeddingLookup(feature='node_feature', ...) to perform embedding lookup on the node_feature component of the GraphTensor, or, EmbeddingLookup(feature='edge_feature', ...) to perform embedding lookup on the edge_feature component of the GraphTensor. If not specified, the node_feature component will be considered.

Instead of specifying feature, NodeEmbeddingLookup(...) or EdgeEmbeddingLookup(...) can be used instead.

Example usage:

>>> graph_tensor = molgraph.GraphTensor(
...     sizes=[2, 3],
...     node_feature=['Sym:C', 'Sym:C', 'Sym:C', 'Sym:O', 'Sym:N'],
...     edge_src=[1, 0, 3, 4, 2, 4, 3, 2],
...     edge_dst=[0, 1, 2, 2, 3, 3, 4, 4],
... )
>>> embedding = molgraph.layers.EmbeddingLookup(
...    feature='node_feature', output_dim=4)
>>> embedding.adapt(graph_tensor)
>>> model = tf.keras.Sequential([embedding])
>>> graph_tensor = model(graph_tensor)
>>> graph_tensor.node_feature.shape
TensorShape([5, 4])

Adapt layer on a tf.data.Dataset:

>>> graph_tensor = molgraph.GraphTensor(
...     sizes=[2, 3],
...     node_feature=['Sym:C', 'Sym:C', 'Sym:C', 'Sym:O', 'Sym:N'],
...     edge_src=[1, 0, 3, 4, 2, 4, 3, 2],
...     edge_dst=[0, 1, 2, 2, 3, 3, 4, 4],
... )
>>> ds = tf.data.Dataset.from_tensor_slices(graph_tensor).batch(2)
>>> embedding = molgraph.layers.EmbeddingLookup(
...    feature='node_feature', output_dim=4)
>>> embedding.adapt(ds)
>>> model = tf.keras.Sequential([embedding])
>>> output = model.predict(ds, verbose=0)
>>> output.node_feature.shape
TensorShape([5, 4])

Parameters

output_dim (int) – The output dimension of the embedding layer.
input_dim (int, None) – The input dimension to the embedding layer. If None, the input dimension is determined by the vocabulary size (obtained from adapting the StringLookup layer to the data). Default to None.
embedding_initializer (tf.keras.initializers.Initializer, str) – Initializer function for the embedding. Default to 'uniform'.
embedding_regularizer (tf.keras.regularizers.Regularizer, None) – Regularizer function applied to the embedding. Default to None.
embedding_constraint (tf.keras.constraints.Constraint, None) – Constraint function applied to the embedding. Default to None.
max_tokens (int, None) – Maximum number of tokens to use. Default to None.
num_oov_indices (int) – Number of out-of-vocabulary indices to use. Default to 1.
mask_token (str) – The token that represents masked input. Default to '[MASK]'.
oov_token (str) – The token that represents out-of-vocabulary input. Default to '[UNK]'.
vocabulary (list, None) – Optional vocabulary. If None, obtain a vocabulary via the adapt() method. Default to None.
**kwargs – Specify the relevant feature. Default to node_feature. The reminaing kwargs are passed to the parent class.

adapt(data, batch_size=None, steps=None)[source]

Adapts the layer to data.

When adapting the layer to the data, build() will be called automatically (to initialize the relevant attributes). After adaption, the layer is finalized and ready to be used.

Parameters

data (GraphTensor, tf.data.Dataset) – Data to be used to adapt the layer. Can be either a GraphTensor directly or a tf.data.Dataset constructed from a GraphTensor.
batch_size (int, None) – The batch size to be used during adaption. Default to None.
steps (int, None) – The number of steps of adaption. If None, the number of samples divided by the batch_size is used. Default to None.

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: data (GraphTensor) – Input to the layer.
Returns: A GraphTensor with updated features. Either the node_feature component or the edge_feature component (of the GraphTensor) are updated.
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.

FeatureProjection

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

Feature projection via dense layer.

Specify, as keyword argument only, FeatureProjection(feature='node_feature') to perform a projection of the node_feature component of the GraphTensor, or, FeatureProjection(feature='edge_feature') to perform a projection of the edge_feature component of the GraphTensor. If not specified, the node_feature component will be considered.

Instead of specifying feature, NodeFeatureProjection(...) or EdgeFeatureProjection(...) can be used instead.

Example usage:

>>> graph_tensor = molgraph.GraphTensor(
...     sizes=[2, 3],
...     node_feature=[[1., 0.], [1., 0.], [1., 0.], [1., 0.], [0., 1.]],
...     edge_src=[1, 0, 3, 4, 2, 4, 3, 2],
...     edge_dst=[0, 1, 2, 2, 3, 3, 4, 4],
... )
>>> model = tf.keras.Sequential([
...     molgraph.layers.FeatureProjection(
...         feature='node_feature', units=16)
... ])
>>> model(graph_tensor).node_feature.shape
TensorShape([5, 16])

Parameters

units (int, None) – Number of output units.
activation (tf.keras.activations.Activation, callable, str, None) – Activation function applied to the output of the layer. Default to None.
use_bias (bool) – Whether the layer should use biases. Default to False.
kernel_initializer (tf.keras.initializers.Initializer, str) – Initializer function for the kernel. Default to tf.keras.initializers.TruncatedNormal(stddev=0.005).
bias_initializer (tf.keras.initializers.Initializer, str) – Initializer function for the bias. Default to tf.keras.initializers.Constant(0.).
kernel_regularizer (tf.keras.regularizers.Regularizer, None) – Regularizer function applied to the kernel. Default to None.
bias_regularizer (tf.keras.regularizers.Regularizer, None) – Regularizer function applied to the bias. 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 kernel. Default to None.
bias_constraint (tf.keras.constraints.Constraint, None) – Constraint function applied to the bias. Default to None.
**kwargs – Specify the relevant feature. Default to node_feature. The reminaing kwargs are passed to the parent class.

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 tf.Tensor or tf.RaggedTensor based on the node_feature component of the inputted 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.

FeatureMasking

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

Randomly masking node or edge features from the graph.

Important: Requires node (or edge) features to be tokenized. I.e., requires the GraphTensor instance to be produced from molgraph.chemistry.Tokenizer instead of molgraph.chemistry.Tokenizer.

Instead of specifying feature, NodeFeatureMasking(...) or EdgeFeatureMasking(...) can be used instead.

Example:

>>> 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)
>>> model = tf.keras.Sequential([
...     molgraph.layers.NodeFeatureMasking(rate=0.15, mask_token='[MASK]'),
...     node_embedding,
...     molgraph.layers.EdgeFeatureMasking(rate=0.15, mask_token='[MASK]'),
...     edge_embedding,
... ])
>>> output = model(graph_tensor)
>>> output.node_feature.shape, output.edge_feature.shape
(TensorShape([5, 32]), TensorShape([8, 32]))

Parameters

rate (float) – The dropout rate. Default to 0.15.
mask_token (str) – The mask token. Default to ‘[MASK]’.
**kwargs – Specify the relevant feature. Default to node_feature. The reminaing kwargs are passed to the parent class.

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 masked instance of GraphTensor.
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.

NodeDropout

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

Randomly dropping nodes from the graph.

Example usage:

>>> 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],
... )
>>> model = tf.keras.Sequential([
...     molgraph.layers.NodeDropout(rate=1.0),
...     molgraph.layers.EdgeDropout(rate=0.0)
... ])
>>> model(graph_tensor).edge_feature.shape
TensorShape([0, 2])

Parameters: rate (float) – The dropout rate. Default to 0.15.

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 instance of GraphTensor with some of its nodes dropped.
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.

EdgeDropout

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

Randomly dropping edges from the graph.

Example usage:

>>> 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],
... )
>>> model = tf.keras.Sequential([
...     molgraph.layers.NodeDropout(rate=0.0),
...     molgraph.layers.EdgeDropout(rate=0.15)
... ])
>>> model(graph_tensor).node_feature.shape
TensorShape([5, 2])

Parameters: rate (float) – The dropout rate. Default to 0.15.

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 instance of GraphTensor with some of its edges dropped.
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.

Postprocessing

DotProductIncident

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

Performs dot product on the incident node features.

Useful for e.g., edge and link classification.

Example usage:

>>> graph_tensor = molgraph.GraphTensor(
...     sizes=[2, 3],
...     node_feature=[[1., 0.], [1., 0.], [1., 0.], [1., 0.], [0., 1.]],
...     edge_src=[1, 0, 3, 4, 2, 4, 3, 2],
...     edge_dst=[0, 1, 2, 2, 3, 3, 4, 4],
... )
>>> model = tf.keras.Sequential([
...     molgraph.layers.DotProductIncident()
... ])
>>> model(graph_tensor)
GraphTensor(
  sizes=<tf.Tensor: shape=(2,), dtype=int32>,
  node_feature=<tf.Tensor: shape=(5, 2), dtype=float32>,
  edge_src=<tf.Tensor: shape=(8,), dtype=int32>,
  edge_dst=<tf.Tensor: shape=(8,), dtype=int32>,
  edge_score=<tf.Tensor: shape=(8, 1), dtype=float32>)

Parameters

normalize (bool) – Whether to apply normalization on the edge scores. Produces cosine similarity values in the range -1 to 1. Default to False.
axes (int, tuple) – The axes (or axis) to perform the dot product. Default to 1.
data_field (str, None) – Name of the data added to the GraphTensor instance. If None, the output will be a tf.Tensor or tf.RaggedTensor containing the dot product between incident node features. If str, a GraphTensor instance with a new data field “data_field” will be outputted. Default to “edge_score”.

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 tf.Tensor or tf.RaggedTensor based on the node_feature field of the inputted 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.

GatherIncident

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

Gathers incident node features.

Useful for e.g., downstream edge and link classification.

Example usage:

>>> graph_tensor = molgraph.GraphTensor(
...     sizes=[2, 3],
...     node_feature=[[1., 0.], [1., 0.], [1., 0.], [1., 0.], [0., 1.]],
...     edge_src=[1, 0, 3, 4, 2, 4, 3, 2],
...     edge_dst=[0, 1, 2, 2, 3, 3, 4, 4],
... )
>>> model = tf.keras.Sequential([
...     molgraph.layers.GatherIncident(concat=True)
... ])
>>> model(graph_tensor)
<tf.Tensor: shape=(8, 4), dtype=float32, numpy=
array([[1., 0., 1., 0.],
       [1., 0., 1., 0.],
       [1., 0., 1., 0.],
       [0., 1., 1., 0.],
       [1., 0., 1., 0.],
       [0., 1., 1., 0.],
       [1., 0., 0., 1.],
       [1., 0., 0., 1.]], dtype=float32)>

Parameters: concat (bool) – Whether to concatenate incident node features or not. If True, resulting shape is (num_edges, num_features * 2), if False, resulting shape is (num_edges, 2, num_features).

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 tf.Tensor or tf.RaggedTensor based on the node_feature field of the inputted 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.

ExtractField

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

Extract specific field of GraphTensor.

Example usage:

>>> graph_tensor = molgraph.GraphTensor(
...     sizes=[2, 3],
...     node_feature=[[1., 0.], [1., 0.], [1., 0.], [1., 0.], [0., 1.]],
...     edge_src=[1, 0, 3, 4, 2, 4, 3, 2],
...     edge_dst=[0, 1, 2, 2, 3, 3, 4, 4],
... )
>>> model = tf.keras.Sequential([
...     molgraph.layers.ExtractField('node_feature')
... ])
>>> model(graph_tensor)
<tf.Tensor: shape=(5, 2), dtype=float32, numpy=
array([[1., 0.],
       [1., 0.],
       [1., 0.],
       [1., 0.],
       [0., 1.]], dtype=float32)>

Parameters: field (str) – Field to extract from GraphTensor.

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 tf.Tensor or tf.RaggedTensor based on the node_feature field of the inputted 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.