matcha.sklearn.graph.gt

Sklearn-compatible Graph Transformer (GT) classifiers and regressors.

Classes

GTClassifier

Graph Transformer (GT) classifier.

GTRegressor

Graph Transformer (GT) regressor.

Module Contents

class matcha.sklearn.graph.gt.GTClassifier(enc_atom_hidden_dim: int = 256, enc_num_layers: int = 3, enc_jk: str = 'concat', enc_num_heads: int = 8, enc_expansion_k: int = 2, enc_distance_k: int | None = 10, enc_readout: str = 'virtualnode', enc_dropout: float = 0.2, enc_activation: str = 'gelu', pred_hidden_dims: list[int] | None = [256, 256], pred_task_head_dims: list[int] | None = None, pred_activation: str = 'gelu', pred_dropout: float = 0.2, num_endpoints: int = 1, loss_fn: str = 'bce', loss_args: dict = {}, optimizer: str = 'adamw', optimizer_args: dict = {'lr': 0.0001, 'weight_decay': 0.0001}, scheduler: str = 'warmup_linear_decay', scheduler_args: dict = {'min_lr': 1e-05}, num_epochs: int = 100, batch_size: int = 64, stochastic_weight_averaging: bool = False, early_stopping: bool = True, patience: int = 20, devices: int = 1, accelerator: str = 'gpu', rwse_k: int = 20, laplacian_k: int = 0, elstatic_k: int = 0, distmat_k: int = 0, rrwp_k: int = 20, num_virtual_nodes: int = 1, feature_list: list[str] | None = None, label_encoder_params: dict = {}, label_transform_map: str | list[str] | dict | None = None, augment_resonance: bool = False, seed: int = 0)[source]

Bases: matcha.sklearn.graph.base_sklearn_gnn.BaseScikitLearnGNN, matcha.sklearn.base_sklearn_model.ScikitLearnClassifierMixin

Graph Transformer (GT) classifier.

A sparse graph transformer that uses multi-head self-attention over graph nodes for molecular property prediction. Compatible with classification datasets only.

Inherits from BaseScikitLearnGNN for graph-specific datamodule creation and collate-aware dataloading.

Loosely inspired by:

Example usage:

model = GTClassifier()
model.fit(train_mols, train_y)
predictions = model.predict(test_mols)
Parameters:

  • enc_atom_hidden_dim (int) – encoder output atom feature dimensionality, defaults to 256

  • enc_num_layers (int) – number of encoder layers, defaults to 3

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

  • enc_num_heads (int) – number of attention heads, defaults to 8

  • enc_expansion_k (int) – FFN expansion factor, defaults to 2

  • enc_distance_k (int | None) – shortest-path distance clipping, defaults to 10

  • enc_readout (str) – encoder aggregation for molecule-level embeddings, defaults to ‘virtualnode’

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

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

  • pred_hidden_dims (list[int] | None) – shape of hidden MLP layers in the predictor, defaults to [256, 256]

  • pred_task_head_dims (list[int] | None) – shape of per-task MLP layers, defaults to None

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

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

  • num_endpoints (int) – number of endpoints or classes to predict, defaults to 1

  • loss_fn (str) – loss function to optimize, defaults to ‘bce’

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

  • optimizer (str) – optimizer to use, defaults to ‘adamw’

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

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

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

  • num_epochs (int) – number of training epochs, defaults to 100

  • batch_size (int) – batch size for training and prediction, defaults to 64

  • stochastic_weight_averaging (bool) – whether to use SWA, defaults to False

  • early_stopping (bool) – whether to use early stopping, defaults to True

  • patience (int) – epochs to wait before early stopping, defaults to 20

  • devices (int) – number of devices for training, defaults to 1

  • accelerator (str) – hardware accelerator, defaults to ‘gpu’

  • rwse_k (int) – random walk structural encoding dimensions, defaults to 20

  • laplacian_k (int) – Laplacian PE components, defaults to 0

  • elstatic_k (int) – electrostatic encoding dimensions, defaults to 0

  • distmat_k (int) – distance matrix encoding dimensions, defaults to 0

  • rrwp_k (int) – relative random walk probability dimensions, defaults to 20

  • num_virtual_nodes (int) – number of virtual nodes, defaults to 1

  • feature_list (list[str] | None) – molecular feature set to compute, defaults to None

  • label_encoder_params (dict) – label encoder parameters, defaults to {}

  • label_transform_map (str | list[str] | dict | None) – label transform specification, defaults to None

  • augment_resonance (bool) – whether to augment with resonance structures, defaults to False

  • seed (int) – random seed, defaults to 0

class matcha.sklearn.graph.gt.GTRegressor(enc_atom_hidden_dim: int = 256, enc_num_layers: int = 3, enc_jk: str = 'concat', enc_num_heads: int = 8, enc_expansion_k: int = 2, enc_distance_k: int | None = 10, enc_readout: str = 'virtualnode', enc_dropout: float = 0.2, enc_activation: str = 'gelu', pred_hidden_dims: list[int] | None = [256, 256], pred_task_head_dims: list[int] | None = None, pred_activation: str = 'gelu', pred_dropout: float = 0.2, num_endpoints: int = 1, loss_fn: str = 'mse', loss_args: dict = {}, optimizer: str = 'adamw', optimizer_args: dict = {'lr': 0.0001, 'weight_decay': 0.0001}, scheduler: str = 'warmup_linear_decay', scheduler_args: dict = {'min_lr': 1e-05}, num_epochs: int = 100, batch_size: int = 64, stochastic_weight_averaging: bool = False, early_stopping: bool = True, patience: int = 20, devices: int = 1, accelerator: str = 'gpu', rwse_k: int = 20, laplacian_k: int = 0, elstatic_k: int = 0, distmat_k: int = 0, rrwp_k: int = 20, num_virtual_nodes: int = 1, feature_list: list[str] | None = None, clip: bool = True, label_encoder_params: dict = {}, label_transform_map: str | list[str] | dict | None = None, scaler_type: str = 'standard', augment_resonance: bool = False, seed: int = 0)[source]

Bases: matcha.sklearn.graph.base_sklearn_gnn.BaseScikitLearnGNN, matcha.sklearn.base_sklearn_model.ScikitLearnRegressorMixin

Graph Transformer (GT) regressor.

A sparse graph transformer that uses multi-head self-attention over graph nodes for molecular property prediction. Compatible with regression datasets only.

Inherits from BaseScikitLearnGNN for graph-specific datamodule creation and collate-aware dataloading.

Loosely inspired by:

Example usage:

model = GTRegressor()
model.fit(train_mols, train_y)
predictions = model.predict(test_mols)
Parameters:

  • enc_atom_hidden_dim (int) – encoder output atom feature dimensionality, defaults to 256

  • enc_num_layers (int) – number of encoder layers, defaults to 3

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

  • enc_num_heads (int) – number of attention heads, defaults to 8

  • enc_expansion_k (int) – FFN expansion factor, defaults to 2

  • enc_distance_k (int | None) – shortest-path distance clipping, defaults to 10

  • enc_readout (str) – encoder aggregation for molecule-level embeddings, defaults to ‘virtualnode’

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

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

  • pred_hidden_dims (list[int] | None) – shape of hidden MLP layers in the predictor, defaults to [256, 256]

  • pred_task_head_dims (list[int] | None) – shape of per-task MLP layers, defaults to None

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

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

  • num_endpoints (int) – number of endpoints to predict, defaults to 1

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

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

  • optimizer (str) – optimizer to use, defaults to ‘adamw’

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

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

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

  • num_epochs (int) – number of training epochs, defaults to 100

  • batch_size (int) – batch size for training and prediction, defaults to 64

  • stochastic_weight_averaging (bool) – whether to use SWA, defaults to False

  • early_stopping (bool) – whether to use early stopping, defaults to True

  • patience (int) – epochs to wait before early stopping, defaults to 20

  • devices (int) – number of devices for training, defaults to 1

  • accelerator (str) – hardware accelerator, defaults to ‘gpu’

  • rwse_k (int) – random walk structural encoding dimensions, defaults to 20

  • laplacian_k (int) – Laplacian PE components, defaults to 0

  • elstatic_k (int) – electrostatic encoding dimensions, defaults to 0

  • distmat_k (int) – distance matrix encoding dimensions, defaults to 0

  • rrwp_k (int) – relative random walk probability dimensions, defaults to 20

  • num_virtual_nodes (int) – number of virtual nodes, defaults to 1

  • feature_list (list[str] | None) – molecular feature set to compute, defaults to None

  • clip (bool) – whether to clip predictions to training label range, defaults to True

  • label_encoder_params (dict) – label encoder parameters, defaults to {}

  • label_transform_map (str | list[str] | dict | None) – label transform specification, defaults to None

  • scaler_type (str) – type of target scaler, defaults to ‘standard’

  • augment_resonance (bool) – whether to augment with resonance structures, defaults to False

  • seed (int) – random seed, defaults to 0