Designing a GCN

Designing a GCN with NAGL primarily involves creating an instance of the GNNModel class. This 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

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

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

model = GNNModel(
    atom_features=atom_features,
    bond_features=bond_features,
    convolution_architecture=gcn.SAGEConvStack,
    n_convolution_hidden_features=128,
    n_convolution_layers=3,
    n_readout_hidden_features=128,
    n_readout_layers=4,
    activation_function=ActivationFunction.ReLU,
    postprocess_layer=postprocess.ComputePartialCharges,
    readout_name=f"am1bcc-charges",
    learning_rate=0.001,
)

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

convolution_architecture: SAGEConv
postprocess_layer: compute_partial_charges

activation_function: ReLU
learning_rate: 0.001
n_convolution_hidden_features: 128
n_convolution_layers: 3
n_readout_hidden_features: 128
n_readout_layers: 4

atom_features:
  - AtomicElement:
      categories: ["C", "H", "O", "N", "P", "S"]
  - AtomConnectivity
  - AtomAverageFormalCharge
  - AtomHybridization
  - AtomInRingOfSize: 3
  - AtomInRingOfSize: 4
  - AtomInRingOfSize: 5
  - AtomInRingOfSize: 6
bond_features:
  - BondOrder
  - BondInRingOfSize: 3
  - BondInRingOfSize: 4
  - BondInRingOfSize: 5
  - BondInRingOfSize: 6

And then loaded with the GNNModel.from_yaml_file() method:

from openff.nagl import GNNModel

model = GNNModel.from_yaml_file("model.yml")

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.nn.gcn module.

Number of Features and Layers

The size and shape of the neural networks in the convolution and readout modules are specified by four arguments:

  • n_convolution_hidden_features

  • n_convolution_layers

  • n_readout_hidden_features

  • n_readout_layers

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.