ChemProp

ChemProp and Chemeleon model implementations.

class openadmet.models.architecture.chemprop.ChemPropModel(*, n_tasks: int = 1, messages: str = 'bond', aggregation: str = 'norm', depth: int = 3, message_hidden_dim: int = 300, ffn_hidden_dim: int = 300, ffn_num_layers: int = 1, normalized_targets: bool = True, batch_norm: bool = False, dropout: float = 0.0, from_foundation: str | None = None, from_chemeleon: bool = False, monitor_metric: str = 'val_loss', metric_list: list = ['mse', 'mae', 'rmse'], scheduler: str = 'noam', max_lr: float = 0.001, weight_decay: float = 0.0, mpnn_lr: float | None = None, ffn_lr: float | None = None, mpnn_weight_decay: float | None = None, ffn_weight_decay: float | None = None, init_lr: float | None = None, final_lr: float | None = None, warmup_epochs: int = 2, reduce_lr_factor: float = 0.1, reduce_lr_patience: int = 10)[source]

Bases: LightningModelBase

ChemProp regression model.

This class implements a ChemProp-based regression model using message passing neural networks (MPNNs) for molecular property prediction. It supports various configurations for message passing, aggregation, and feed-forward network (FFN) layers. Can be initialized from the CheMeleon foundation model [REF], overriding settings for depth, message hidden dim, messages, and aggregation.

Variables:

  • type (str) – The type of the model.

  • n_tasks (int) – Number of prediction tasks.

  • messages (str) – Type of message passing (“bond” or “atom”).

  • aggregation (str) – Aggregation method (“mean” or “norm”).

  • depth (int) – Number of message passing steps.

  • message_hidden_dim (int) – Hidden dimension size for message passing.

  • ffn_hidden_dim (int) – Hidden dimension size for the feed-forward network.

  • ffn_num_layers (int) – Number of layers in the feed-forward network.

  • normalized_targets (bool) – Whether targets are normalized.

  • batch_norm (bool) – Whether to use batch normalization.

  • dropout (float) – Dropout rate.

  • from_chemeleon (bool) – Whether to use the CheMeleon foundation model. Deprecated — use from_foundation='chemeleon' instead.

  • monitor_metric (str) – The metric to monitor during training. Default is “val_loss”.

  • metric_list (list) – List of metrics to use for evaluation. Default is [“mse”, “mae”, “rmse”].

  • scheduler (str) – Learning rate scheduler (“noam” or “plateau”). Default is “noam”.

  • warmup_epochs (int) – Number of warmup epochs for learning rate scheduling (Noam scheduler only). Default is 2.

  • init_lr (float, optional) – Initial learning rate. If None, defaults to max_lr * 0.1.

  • max_lr (float) – Maximum learning rate (Global default). Default is 1e-3.

  • final_lr (float, optional) – Final learning rate. If None, defaults to max_lr * 0.01.

  • weight_decay (float) – Global weight decay. Default is 0.0.

  • mpnn_lr (float, optional) – Learning rate for MPNN. If None, defaults to max_lr.

  • ffn_lr (float, optional) – Learning rate for FFN. If None, defaults to max_lr.

  • mpnn_weight_decay (float, optional) – Weight decay for MPNN. If None, defaults to weight_decay.

  • ffn_weight_decay (float, optional) – Weight decay for FFN. If None, defaults to weight_decay.

  • reduce_lr_factor (float) – Factor by which the learning rate will be reduced (Plateau scheduler only). Default is 0.1.

  • reduce_lr_patience (int) – Number of epochs with no improvement after which learning rate will be reduced (Plateau scheduler only). Default is 10.

aggregation: str
batch_norm: bool
build(scaler=None)[source]

Prepare and build the ChemProp model.

Downloads and loads the CheMeleon foundation model if specified, otherwise constructs a new MPNN model with the given configuration.

Parameters:

scaler (object, optional) – Scaler for target normalization.

Returns:

self – The current instance with the estimator built.

Return type:

ChemPropModel

depth: int
dropout: float
ffn_hidden_dim: int
ffn_lr: float | None
ffn_num_layers: int
ffn_weight_decay: float | None
final_lr: float | None
freeze_weights(message_passing: bool = True, batch_norm: bool = True, ffn_layers: int = 0)[source]

Freeze parts of the model for transfer learning or fine-tuning.

Parameters:

  • message_passing (bool, optional) – If True, freeze the message passing layers. Default is True.

  • batch_norm (bool, optional) – If True, freeze the batch normalization layers. Default is True.

  • ffn_layers (int, optional) – Number of feed-forward network (FFN) layers to freeze. Default is 0.

Notes

This method sets the requires_grad attribute of the specified layers to False, preventing their weights from being updated during training. It also sets these layers to evaluation mode.

from_chemeleon: bool
from_foundation: str | None
init_lr: float | None
make_new() → ChemPropModel[source]

Copy parameters to a new model instance without copying the estimator.

max_lr: float
message_hidden_dim: int
messages: str
metric_list: list
model_config: ClassVar[ConfigDict] = {}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

model_post_init(context: Any, /) → None

This function is meant to behave like a BaseModel method to initialize private attributes.

It takes context as an argument since that’s what pydantic-core passes when calling it.

Args:

self: The BaseModel instance. context: The context.

monitor_metric: str
mpnn_lr: float | None
mpnn_weight_decay: float | None
n_tasks: int
normalized_targets: bool
predict(X: ndarray, accelerator='gpu', devices=1, **kwargs) → ndarray[source]

Predict using the trained model.

Parameters:

  • X (np.ndarray) – Input data for prediction.

  • accelerator (str, optional) – Accelerator type to use (“gpu” or “cpu”).

  • devices (int, optional) – Number of devices to use for prediction.

  • **kwargs – Additional keyword arguments for the trainer.

Returns:

Model predictions.

Return type:

np.ndarray

reduce_lr_factor: float
reduce_lr_patience: int
resolve_hyperparameters() → ChemPropModel[source]

Resolve hyperparameters using global defaults and component overrides pattern.

Logic: - Resolve learning rates:

  • init_lr -> max_lr * 0.1

  • final_lr -> max_lr * 0.01

  • mpnn_lr -> max_lr

  • ffn_lr -> max_lr

  • Resolve weight decays:

    • mpnn_weight_decay -> weight_decay

    • ffn_weight_decay -> weight_decay

scheduler: str
serialize(param_path='model.json', serial_path='model.pth')[source]

Save the model with only explicitly provided parameter fields.

Parameters:

  • param_path (PathLike) – Path to save the model parameters to

  • serial_path (PathLike) – Path to save the serialized model to

set_n_tasks() → ChemPropModel[source]

Set the number of tasks for the model.

Returns:

The updated model instance.

Return type:

ChemPropModel

train(dataloader, scaler=None)[source]

Train the model.

Parameters:

  • dataloader (DataLoader) – DataLoader for training data.

  • scaler (object, optional) – Scaler for target normalization.

type: ClassVar[str] = 'ChemPropModel'
validate_scheduler_params() → ChemPropModel[source]

Ensure scheduler-specific parameters are valid for the chosen scheduler.

warmup_epochs: int
weight_decay: float
openadmet.models.architecture.chemprop.configure_optimizers(self)[source]

Configure optimizers and learning rate schedulers.

Returns:

A dictionary containing the optimizer and learning rate scheduler configurations.

Return type:

dict