molgraph 0.6.2

Implementations of graph neural networks for molecular machine learning

MolGraph

Graph Neural Networks with TensorFlow and Keras. Focused on Molecular Machine Learning.

This is an early release and a work in progress; so be aware of breaking changes.

Highlights

Build a Graph Neural Network with Keras' Sequential API:

from molgraph import GraphTensor
from molgraph import layers
from tensorflow import keras

model = keras.Sequential([
    layers.GINConv(units=32),
    layers.GINConv(units=32),
    layers.Readout(),
    keras.layers.Dense(units=1),
])
output = model(
    GraphTensor(node_feature=[[4.], [2.]], edge_src=[0], edge_dst=[1])
)

Paper

See arXiv

Documentation

See readthedocs

Implementations

• Graph tensor (GraphTensor)
  • A composite tensor holding graph data.
  • Has a ragged (multiple graphs) and a non-ragged state (single disjoint graph)
  • Can conveniently go between both states (merge(), separate())
  • Can propagate node information (features) based on edges (propagate())
  • Can add, update and remove graph data (update(), remove())
  • As it is now implemented with the TF's ExtensionType API, it is now compatible with TensorFlow's APIs (including Keras). For instance, graph data (encoded as a GraphTensor) can now seamlessly be used with keras.Sequential, keras.Functional, tf.data.Dataset, and tf.saved_model APIs.
• Layers
  • Convolutional
    • GCNConv (GCNConv)
    • GINConv (GINConv)
    • GCNIIConv (GCNIIConv)
    • GraphSageConv (GraphSageConv)
  • Attentional
    • GATConv (GATConv)
    • GATv2Conv (GATv2Conv)
    • GTConv (GTConv)
    • GMMConv (GMMConv)
    • GatedGCNConv (GatedGCNConv)
    • AttentiveFPConv (AttentiveFPConv)
  • Message-passing
    • MPNNConv (MPNNConv)
    • EdgeConv (EdgeConv)
  • Distance-geometric
    • DTNNConv (DTNNConv)
    • GCFConv (GCFConv)
  • Pre- and post-processing
    • In addition to the aforementioned GNN layers, there are also several other layers which improves model-building. See readout/, preprocessing/, postprocessing/, positional_encoding/.
• Models
  • Although model building is easy with MolGraph, there are some built-in GNN models:
    • GIN
    • MPNN
    • DMPNN
  • And models for improved interpretability of GNNs:
    • SaliencyMapping
    • IntegratedSaliencyMapping
    • SmoothGradSaliencyMapping
    • GradientActivationMapping (Recommended)

Changelog

For a detailed list of changes, see the CHANGELOG.md.

Requirements/dependencies

• Python (version >= 3.10 recommended)
  • TensorFlow (version >= 2.13.0 recommended)
  • RDKit (version >= 2022.3.5 recommended)
  • Pandas (version >= 1.0.3 recommended)
  • IPython (version == 8.12.0 recommended)

MolGraph should work with the more recent TensorFlow and RDKit versions. If not, try installing earlier versions of TensorFlow and RDKit.

Installation

Install via pip:

pip install molgraph

Install via docker:

git clone https://github.com/akensert/molgraph.git
cd molgraph/docker
docker build -t molgraph-tf[-gpu][-jupyter]/molgraph:0.0 molgraph-tf[-gpu][-jupyter]/
docker run -it [-p 8888:8888] molgraph-tf[-gpu][-jupyter]/molgraph:0.0

Now run your first program with MolGraph:

from tensorflow import keras
from molgraph import chemistry
from molgraph import layers
from molgraph import models

# Obtain dataset, specifically ESOL
qm7 = chemistry.datasets.get('esol')

# Define molecular graph encoder
atom_encoder = chemistry.Featurizer([
    chemistry.features.Symbol(),
    chemistry.features.Hybridization(),
    # ...
])

bond_encoder = chemistry.Featurizer([
    chemistry.features.BondType(),
    # ...
])

encoder = chemistry.MolecularGraphEncoder(atom_encoder, bond_encoder)

# Obtain graphs and associated labels
x_train = encoder(qm7['train']['x'])
y_train = qm7['train']['y']

x_test = encoder(qm7['test']['x'])
y_test = qm7['test']['y']

# Build model via Keras API
gnn_model = keras.Sequential([
    layers.GATConv(units=32, name='gat_conv_1'),
    layers.GATConv(units=32, name='gat_conv_2'),
    layers.Readout(),
    keras.layers.Dense(units=1024, activation='relu'),
    keras.layers.Dense(units=y_train.shape[-1])
])

# Compile, fit and evaluate
gnn_model.compile(optimizer='adam', loss='mae')
gnn_model.fit(x_train, y_train, epochs=50)
scores = gnn_model.evaluate(x_test, y_test)

# Compute gradient activation maps
gam_model = models.GradientActivationMapping(
    model=gnn_model, layer_names=['gat_conv_1', 'gat_conv_2'])

maps = gam_model(x_train.separate())