Designing a GCN

Designing a GCN with NAGL primarily involves creating an instance of the [ModelConfig] class.

In Python

A ModelConfig class can be created can be done straightforwardly in Python.

from openff.nagl.features import atoms, bonds
from openff.nagl import GNNModel
from openff.nagl.nn import gcn
from openff.nagl.nn.activation import ActivationFunction
from openff.nagl.nn import postprocess

from openff.nagl.config.model import (
    ModelConfig,
    ConvolutionModule, ReadoutModule,
    ConvolutionLayer, ForwardLayer,
)

First we can specify our desired features. These should be instances of the feature classes.

atom_features = (
    atoms.AtomicElement(["C", "H", "O", "N", "P", "S"]),
    atoms.AtomConnectivity(),
    ...
)

bond_features = (
    bonds.BondOrder(),
    ...
)

Next, we can design convolution and readout modules. For example:

convolution_module = ConvolutionModule(
    architecture="SAGEConv",
    # construct 2 layers with dropout 0 (default),
    # hidden feature size 512, and ReLU activation function
    # these layers can also be individually specified,
    # but we just duplicate the layer 6 times for identical layers
    layers=[
        ConvolutionLayer(
            hidden_feature_size=512,
            activation_function="ReLU",
            aggregator_type="mean"
        )
    ] * 2,
)

# define our readout module/s
# multiple are allowed but let's focus on charges
readout_modules = {
    # key is the name of output property, any naming is allowed
    "charges": ReadoutModule(
        pooling="atoms",
        postprocess="compute_partial_charges",
        # 1 layers
        layers=[
            ForwardLayer(
                hidden_feature_size=512,
                activation_function="ReLU",
            )
        ],
    )
}

We can now bring it all together as a ModelConfig and create a GNNModel.

model_config = ModelConfig(
    version="0.1",
    atom_features=atom_features,
    bond_features=bond_features,
    convolution=convolution_module,
    readouts=readout_modules,
)

model = GNNModel(model_config)

From YAML

Or if you prefer, the same model architecture can be specified as a YAML file:

version: '0.1'
convolution:
  architecture: SAGEConv
  layers:
    - hidden_feature_size: 512
      activation_function: ReLU
      dropout: 0
      aggregator_type: mean
    - hidden_feature_size: 512
      activation_function: ReLU
      dropout: 0
      aggregator_type: mean
readouts:
  charges:
    pooling: atoms
    postprocess: compute_partial_charges
    layers:
      - hidden_feature_size: 128
        activation_function: Sigmoid
        dropout: 0
atom_features:
  - name: atomic_element
    categories: ["C", "H", "O", "N", "P", "S"]
  - name: atom_connectivity
    categories: [1, 2, 3, 4, 5, 6]
  - name: atom_hybridization
  - name: atom_in_ring_of_size
    ring_size: 3
  - name: atom_in_ring_of_size
    ring_size: 4
  - name: atom_in_ring_of_size
    ring_size: 5
  - name: atom_in_ring_of_size
    ring_size: 6
bond_features:
  - name: bond_is_in_ring

And then loaded into a config using the [ModelConfig.from_yaml()] method:

from openff.nagl import GNNModel
from openff.nagl.config import ModelConfig

model = GNNModel(ModelConfig.from_yaml("model.yaml"))

Here we’ll go through each option, what it means, and where to find the available choices.

atom_features and bond_features

These arguments specify the featurization scheme for the model (see Featurization). atom_features takes a tuple of features from the openff.nagl.features.atoms module, and bond_features a tuple of features from the openff.nagl.features.bonds module. Each feature is a class that must be instantiated, possibly with some arguments. Custom features may be implemented by subclassing AtomFeature or BondFeature; both share the interface of their base class Feature.

convolution_architecture

The convolution_architecture argument specifies the structure of the convolution module. Available options are provided in the [openff.nagl.config.model] module.

Number of Features and Layers

Each module comprises a number of layers that must be individually specified. For example, a [ConvolutionModule] consists of specified [ConvolutionLayer]s. A [ReadoutModule] consists of specified [ForwardLayer]s.

The “convolution” arguments define the update network in the convolution module, and the “readout” the network in the readout module (see The Convolution Module: Message-passing Graph Convolutional Networks and The Readout Module: The Charge Equilibration Method). Read the GNNModel docstring carefully to determine which layers are considered hidden.

activation_function

The activation_function argument defines the activation function used by the readout network (see Neural Networks - a quick primer). The activation function used by the convolution network is currently fixed to ReLU. Available activation functions are the variants (attributes) of the ActivationFunction class. When using the Python API rather than the YAML interface, other activation functions from PyTorch can be used.

postprocess_layer

postprocess_layer specifies a PyTorch Module that performs post-processing in the readout module (see The Readout Module: The Charge Equilibration Method). Post-processing layers inherit from the PostprocessLayer class and are found in the openff.nagl.nn.postprocess module.