Graphium: Scaling Molecular GNNs to Infinity

The graphium/utils directory contains helper functions and utility modules used throughout the library. Key components include:

• spaces.py: Used for tracking the namespace/space of names and dictionaries.
• mup.py: Provides helper functions for MUP (Matrix Unit Product or similar internal mathematical operation) operations.

Graphium uses a configuration system based on Hydra to manage experiment settings. The configuration logic is split into several components:

• _load.py and config_convert.py: Handle the initial loading of settings from yaml configuration files.
• _loader.py: Responsible for retrieving specific arguments from the loaded configuration and orchestrating the loading of different modules with their corresponding arguments.

The graphium/data directory contains the core components for managing data lifecycles in Graphium. It provides tools for loading data from disk, processing it into PyTorch DataLoaders, and defining custom datasets and transformations.

Key components include:

• Data Loading & Processing: datamodule.py handles loading from disk and processing into a DataLoader; utils.py provides utility functions for dataset loading.
• Dataset & Collation: dataset.py defines the dataset class; collate.py defines the collation function used by the DataLoader.
• Transformations & Normalization: smiles_transform.py transforms SMILES strings into molecule objects; normalization.py handles label normalization.
• Specialized Converters: sdf2csv.py provides conversion functions from SDF to CSV specifically for the OGB PCQM dataset.

The graphium.nn module provides the foundational base structures for building neural networks within the Graphium library. These classes are designed to be inherited by other modules to implement custom graph-based neural network architectures. The module is organized into three main components:

1. Base Graph Layer: The fundamental building block for graph-specific layers.
2. Base Layers: Standard neural network layers adapted for graph operations.
3. Residual Connections: Implementations for adding skip connections to facilitate training deeper networks.

The graphium/features directory contains the core logic for transforming molecular data into graph-based representations. Key components include:

• Molecule Featurization: featurizer.py provides the logic to add node, edge, and graph-level features to molecular objects.
• Positional and Structural Encodings:
  • positional_encoding.py: Computes raw positional and structural encodings via the graph_positional_encoder function.
  • rw.py: Computes Random Walk (RW) positional encodings.
  • spectral.py: Computes spectral positional encodings, such as Laplacian eigenvalues and eigenvectors.
• Molecular Properties: properties.py contains functions for computing various molecular properties.
• Utilities: nmp.py provides helper functions (e.g., checking if a string can be converted to a float) used during the featurization process.

The graphium.nn.encoders module provides various implementations of positional encoders used to inject structural or spatial information into graph neural networks.

Available encoder types include:

• Base Encoder: The foundational interface for all encoders.
• Gaussian Kernal Positional Encoder: Uses Gaussian kernels for positional encoding.
• Laplacian Positional Encoder: Uses the Laplacian spectrum for encoding.
• MLP Encoder: Uses a Multi-Layer Perceptron to learn positional representations.
• Signnet Positional Encoder: Implements Signnet-based positional encoding.

The graphium.finetuning module provides tools for finetuning models and performing linear probing (also known as fingerprinting).

Key components include:

• GraphFinetuning: The primary class for managing the finetuning process.
• FullGraphFinetuningNetwork: An architecture for full network finetuning.
• PretrainedModel: A wrapper or representation for the base model being finetuned.
• FinetuningHead: The specific layer or module added to the pretrained model for the finetuning task.
• Fingerprinter: A specialized tool for performing fingerprinting (linear probing) tasks.

The graphium/nn/pyg_layers/ directory contains implementations of various Graph Neural Network (GNN) and Graph Transformer layers designed to work with PyTorch Geometric (PyG).

Available layer implementations include:

• PNA: PNAMessagePassingPyg class.
• GPS (Graph Transformer): GPSLayerPyg class.
• Gated GCN: GatedGCNPyg class.
• GIN/GINE: GINConvPyg and GINEConvPyg classes.
• MPNN: MPNNPlusPyg class.
• Pooling: Various pooling layers provided via pooling_pyg.py.

The graphium command is the primary entrypoint for managing datasets and running fine-tuning utilities. It follows a sub-command structure.

General Usage:

$ graphium [OPTIONS] COMMAND [ARGS]...

Available Commands:

• data: Manage Graphium datasets (download, list, prepare).
• finetune: Access fine-tuning utilities like ADMET fine-tuning and fingerprint extraction.