ChemProp

ChemProp and Chemeleon model implementations.

class openadmet.models.architecture.chemprop.ChemPropModel(*, n_tasks: int = 1, messages: str = 'bond', aggregation: str = 'mean', depth: int = 3, message_hidden_dim: int = 300, ffn_hidden_dim: int = 300, ffn_num_layers: int = 2, 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 | None = None, reduce_lr_factor: float | None = None, reduce_lr_patience: int | None = None, monitor_metric_mode: str = 'min')[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”). Default is “mean”, matching the original ChemProp paper baseline. “norm” uses a learned normalization parameter instead.

  • 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. Default is 2 (one hidden layer); setting to 1 reduces the FFN to a single linear readout, making ffn_hidden_dim irrelevant.

  • 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”.

    Selection depends on the training regime:

    • ”noam”: for fixed-length training where the epoch budget is known in advance. The learning rate follows a preset trajectory, a linear ramp followed by a smooth decay across the configured run. This is the original ChemProp recipe and the appropriate default for standard from-scratch training. It depends on max_epochs being set correctly; an incorrect or open-ended budget distorts the trajectory.

    • ”plateau”: for runs whose length is not fixed in advance, such as early-stopped training or fine-tuning from a foundation model. The learning rate is reduced only when the monitored metric stops improving, adapting to observed progress rather than a preset timeline. Requires a validation set.

    For open-ended training (max_epochs=-1) or early stopping, prefer “plateau”; “noam” cannot shape its trajectory without a known budget.

  • max_lr (float) – Peak learning rate (global reference). Default is 1e-3.

  • final_lr (float, optional) – Floor LR for each param group. Defaults to max_lr * 0.01. When mpnn_lr or ffn_lr differ from max_lr, the absolute floor for that group is group_lr * (final_lr / max_lr); the ratio is preserved proportionally.

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

  • mpnn_lr (float, optional) – Peak learning rate for the MPNN param group. If None, defaults to max_lr.

  • ffn_lr (float, optional) – Peak learning rate for the FFN param group. If None, defaults to max_lr.

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

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

  • warmup_epochs (int, optional) – [Noam only] Number of linear-ramp epochs before geometric decay. If None (default), resolves to 2 (matching the original ChemProp paper default). Setting this field with scheduler=”plateau” raises ValueError. The schedule shape depends on max_epochs being set correctly in the Lightning Trainer; leaving it at the Lightning default (1000) when actual training runs shorter will under-decay the LR.

  • init_lr (float, optional) – [Noam only] Starting LR at the beginning of the warmup ramp. Defaults to max_lr * 0.1. When mpnn_lr or ffn_lr differ from max_lr, the absolute starting LR for that group is group_lr * (init_lr / max_lr); the schedule shape is preserved proportionally around each group’s peak.

  • reduce_lr_factor (float, optional) – [Plateau only] Multiplicative factor applied when a plateau is detected. If None (default), resolves to 0.5. Must be < 1.0. Setting with scheduler=”noam” raises ValueError.

  • reduce_lr_patience (int, optional) – [Plateau only] Epochs with no improvement before LR is reduced. If None (default), resolves to 5. Setting with scheduler=”noam” raises ValueError.

  • monitor_metric_mode (str) – Direction for metric monitoring: “min” for loss-style metrics, “max” for score-style metrics. Default is “min”. Currently consumed by the plateau scheduler; must also match any early-stopping callback monitoring the same metric.

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
monitor_metric_mode: 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 | None
reduce_lr_patience: int | None
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

  • Fill scheduler-specific defaults (only for the active scheduler):

    • noam: warmup_epochs -> 2

    • plateau: reduce_lr_factor -> 0.5, reduce_lr_patience -> 5

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

Save the model with explicitly provided fields plus resolved LR hyperparameters.

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.

Cross-scheduler params use None as the “not set” sentinel so this validator can distinguish user-supplied values from unset fields without relying on model_fields_set (which only tracks explicitly provided keys).

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

Configure optimizers and learning rate schedulers.

Returns:

A dictionary containing the optimizer and learning rate scheduler configurations.

Return type:

dict