Molgraph

<img src="https://github.com/akensert/molgraph/blob/main/docs/source/_static/molgraph-logo-pixel.png" alt="molgraph-title" width="90%">

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

[!NOTE] For compatability with Keras 3, see the more recent project MolCraft. MolCraft also provides improved featurization of molecules (including the addition of super nodes) and improved modules (including models.GraphModel, layers.GraphLayer and tensors.GraphTensor).

Quick start

Benchmark the performance of MolGraph here, and implement a complete model pipeline with MolGraph here.

Highlights

Build a Graph Neural Network with Keras' Sequential API:

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

g = GraphTensor(node_feature=[[4.], [2.]], edge_src=[0], edge_dst=[1])

model = keras.Sequential([
    layers.GNNInput(type_spec=g.spec),
    layers.GATv2Conv(units=32),
    layers.GATv2Conv(units=32),
    layers.Readout(),
    keras.layers.Dense(units=1),
])

pred = model(g)

# Save and load Keras model
model.save('/tmp/gatv2_model.keras')
loaded_model = keras.models.load_model('/tmp/gatv2_model.keras')
loaded_pred = loaded_model(g)
assert pred == loaded_pred

Combine outputs of GNN layers to improve predictive performance:

model = keras.Sequential([
    layers.GNNInput(type_spec=g.spec),
    layers.GNN([
        layers.FeatureProjection(units=32),
        layers.GINConv(units=32),
        layers.GINConv(units=32),
        layers.GINConv(units=32),
    ]),
    layers.Readout(),
    keras.layers.Dense(units=128),
    keras.layers.Dense(units=1),
])

model.summary()

Installation

For CPU users:

<pre> pip install molgraph </pre>

For GPU users:

<pre> pip install molgraph[gpu] </pre>

Implementations

• Tensors
  • Graph tensor
    • A composite tensor holding graph data; compatible with tf.data.Dataset, keras.Sequential and much more.
• Layers
  • Graph convolutional layers
  • Graph attentional layers
  • Graph message passing layers
  • Graph readout layers
  • Preprocessing layers
  • Postprocessing layers
  • Positional encoding layers
• Models
  • Graph neural networks
  • Saliency mapping

Overview

<img src="https://github.com/akensert/molgraph/blob/main/docs/source/_static/molgraph-overview.png" alt="molgraph-overview" width="90%">

Documentation

See readthedocs

Papers

• MolGraph: a Python package for the implementation of molecular graphs and graph neural networks with TensorFlow and Keras. Pre-print here.
• A hands-on tutorial on quantitative structure-activity relationships using fully expressive graph neural networks.