matcha.torch.models.classic.e3gnn_model

E(3)-equivariant graph neural network (E3GNN) classic model.

Classes

E3GNNModel

E(3)-equivariant graph neural network (E3GNN) for molecular property

Module Contents

class matcha.torch.models.classic.e3gnn_model.E3GNNModel(additional_mol_features_dim: int = 0, enc_num_layers: int = 3, enc_atom_input_dim: int = 44, enc_bond_input_dim: int = 14, enc_atom_hidden_dim: int = 128, enc_m_dim: int = 16, enc_fourier_features: int = 4, enc_soft_edge: bool = False, enc_norm_feats: bool = True, enc_norm_coors: bool = True, enc_update_coors: bool = True, enc_jk: str = 'last', enc_readout: str = 'vpa', enc_activation: str = 'swish', enc_dropout: float = 0.2, enc_laplacian_k: int = 10, enc_rwse_k: int = 20, enc_elstatic_k: int = 0, enc_distmat_k: int = 0, enc_rrwp_k: int = 0, pred_hidden_dims: list[int] = [512, 256], pred_task_head_dims: list[int] | None = None, pred_activation: str = 'swish', pred_dropout: float = 0.2, num_endpoints: int = 1, loss_fn: str = 'mse', loss_args: dict = {}, optimizer: str = 'adam', optimizer_args: dict = {'lr': 0.001}, scheduler: str = 'cosine_annealing', scheduler_args: dict = {'min_lr': 1e-06, 'total_steps': 50})[source]

Bases: matcha.torch.models.classic.base_classic_model.BaseClassicModel, lightning.pytorch.core.mixins.HyperparametersMixin

E(3)-equivariant graph neural network (E3GNN) for molecular property prediction from 3-D conformers.

Uses E(n)-equivariant message passing that operates on both node features and 3-D coordinates. Inherits from BaseClassicModel for common training/prediction routines and from HyperparametersMixin for saving hyperparameters.

Reference: Satorras et al., E(n) Equivariant Graph Neural Networks (https://arxiv.org/abs/2102.09844)

Example usage:

model = E3GNNModel()
trainer = L.Trainer(max_epochs=50)
trainer.fit(model=model, train_dataloaders=train_dataloader)
Parameters:

  • additional_mol_features_dim (int) – dimensionality of extra molecular features concatenated to encoder output, defaults to 0

  • enc_num_layers (int) – number of equivariant message-passing layers, defaults to 3

  • enc_atom_input_dim (int) – input atom feature dimensionality, defaults to 44

  • enc_bond_input_dim (int) – input bond feature dimensionality, defaults to 14

  • enc_atom_hidden_dim (int) – hidden atom feature dimensionality, defaults to 128

  • enc_m_dim (int) – message dimensionality, defaults to 16

  • enc_fourier_features (int) – number of Fourier features for distance encoding, defaults to 4

  • enc_soft_edge (bool) – whether to use soft edge weights, defaults to False

  • enc_norm_feats (bool) – whether to normalise node features, defaults to True

  • enc_norm_coors (bool) – whether to normalise coordinate updates, defaults to True

  • enc_update_coors (bool) – whether to update coordinates, defaults to True

  • enc_jk (str) – jumping knowledge strategy, defaults to ‘last’

  • enc_readout (str) – graph-level readout strategy, defaults to ‘vpa’

  • enc_activation (str) – activation function in the encoder, defaults to ‘swish’

  • enc_dropout (float) – dropout rate in the encoder, defaults to 0.2

  • enc_laplacian_k (int) – Laplacian positional encoding dimension, defaults to 10

  • enc_rwse_k (int) – random-walk structural encoding dimension, defaults to 20

  • enc_elstatic_k (int) – electrostatic encoding dimension, defaults to 0

  • enc_distmat_k (int) – distance matrix encoding dimension, defaults to 0

  • enc_rrwp_k (int) – relative random-walk probabilities dimension, defaults to 0

  • pred_hidden_dims (list[int]) – hidden layer sizes in the MLP predictor, defaults to [512, 256]

  • pred_task_head_dims (list[int] | None) – per-task head dimensions, defaults to None

  • pred_activation (str) – activation in the predictor, defaults to ‘swish’

  • pred_dropout (float) – dropout rate in the predictor, defaults to 0.2

  • num_endpoints (int) – number of prediction targets, defaults to 1

  • loss_fn (str) – loss function name, defaults to ‘mse’

  • loss_args (dict) – additional loss function arguments, defaults to {}

  • optimizer (str) – optimizer name, defaults to ‘adam’

  • optimizer_args (dict) – additional optimizer arguments, defaults to {‘lr’: 1e-3}

  • scheduler (str) – learning rate scheduler name, defaults to ‘cosine_annealing’

  • scheduler_args (dict) – additional scheduler arguments, defaults to {‘min_lr’: 1e-6, ‘total_steps’: 50}

params
encoder