MoamaMolecularEncoder¶

Motif-aware attribute masking and fingerprint reconstruction.

from torch_molecule import MoamaMolecularEncoder

Examples · Methods

Constructor¶

class MoamaMolecularEncoder(mask_rate=0.15, lw_rec=0.5, num_layer=5, hidden_size=300, drop_ratio=0.5, norm_layer='batch_norm', encoder_type='gin-virtual', readout='sum', batch_size=128, epochs=500, learning_rate=0.001, grad_clip_value=None, weight_decay=0.0, use_lr_scheduler=False, scheduler_factor=0.5, scheduler_patience=5, *, device=None, model_name='MoamaMolecularEncoder', verbose='none')[source]¶

This encoder implements Motif-aware Attribute Masking for Molecular Graph Pre-training.

References

Paper: https://openreview.net/forum?id=uqPnesiGGi
Code: https://github.com/einae-nd/MoAMa-dev

Parameters:

mask_rate (float, default=0.15) – Fraction of nodes to mask in each graph
lw_rec (float, default=0.5) – Weight balancing between reconstruction loss and fingerprint loss. Higher values emphasize reconstruction, lower values emphasize fingerprint matching.
num_layer (int, default=5) – Number of GNN layers.
hidden_size (int, default=300) – Dimension of hidden node features.
drop_ratio (float, default=0.5) – Dropout probability.
norm_layer (str, default="batch_norm") – Type of normalization layer to use. One of [“batch_norm”, “layer_norm”, “instance_norm”, “graph_norm”, “size_norm”, “pair_norm”].
encoder_type (str, default="gin-virtual") – Type of GNN architecture to use. One of [“gin-virtual”, “gcn-virtual”, “gin”, “gcn”].
readout (str, default="sum") – Method for aggregating node features to obtain graph-level representations. One of [“sum”, “mean”, “max”].
batch_size (int, default=128) – Number of samples per batch for training.
epochs (int, default=500) – Maximum number of training epochs.
learning_rate (float, default=0.001) – Learning rate for optimizer.
grad_clip_value (float, optional) – Maximum norm of gradients for gradient clipping.
weight_decay (float, default=0.0) – L2 regularization strength.
use_lr_scheduler (bool, default=False) – Whether to use a learning rate scheduler.
scheduler_factor (float, default=0.5) – Factor by which to reduce the learning rate when plateau is detected.
scheduler_patience (int, default=5) – Number of epochs with no improvement after which learning rate will be reduced.
verbose (str, default="none") – Whether to display progress info. Options are: “none”, “progress_bar”, “print_statement”. If any other, “none” is automatically chosen.
device (Optional[Union[torch.device, str]], default=None) – Device to run the model on (CPU or GPU).
model_name (str, default="MoamaMolecularEncoder") – Name identifier for the model.

Example¶

examples/encoder/moama.py¶

import numpy as np
from torch_molecule import MoamaMolecularEncoder
from torch_molecule.utils.search import ParameterType, ParameterSpec

def test_moama_encoder():
    # Test molecules (simple examples)
    molecules = [
        "CC(=O)O",  # Acetic acid
        "CCO",      # Ethanol
        "CCCC",     # Butane
        "c1ccccc1", # Benzene
        "CCN",      # Ethylamine
    ]

    # Basic initilization test
    model = MoamaMolecularEncoder(
        num_layer=3,
        hidden_size=300,
        batch_size=5,
        epochs=5,  # Small number for testing
        verbose='print_statement'
    )
    print("Model initialized successfully")
    
    # Basic self-supervised fitting test
    print("\n=== Testing MoAMa model self-supervised fitting ===")
    model.fit(molecules[:4])

    # test predict
    print("\n=== Testing MoAMa model prediction ===")
    vectors = model.encode(molecules[4:])
    print(f"Representation shape: {vectors.shape}")
    print(f"Representation for new molecule: {vectors[0]}")
    # Model saving and loading test
    print("\n=== Testing model saving and loading ===")
    save_path = "test_model.pt"
    model.save_to_local(save_path)
    print(f"Model saved to {save_path}")

    new_model = MoamaMolecularEncoder()
    new_model.load_from_local(save_path)
    print("Model loaded successfully")
    
    # Clean up
    import os
    if os.path.exists(save_path):
        os.remove(save_path)
        print(f"Cleaned up {save_path}")

def test_moama_encoder_polymers():
    # Test molecules (simple examples)
    polymers = [
        "*Nc1ccc([C@H](CCC)c2ccc(C3(c4ccc([C@@H](CCC)c5ccc(N*)cc5)cc4)CCC(CCCCC)CC3)cc2)cc1",
        "*Nc1ccc(-c2c(-c3ccc(C)cc3)c(-c3ccc(C)cc3)c(N*)c(-c3ccc(C)cc3)c2-c2ccc(C)cc2)cc1",
        "*CC(*)(C)C(=O)OCCCCCCCCCOc1ccc2cc(C(=O)Oc3ccccc3)ccc2c1"
    ]
    model = MoamaMolecularEncoder(
        num_layer=3,
        hidden_size=300,
        batch_size=5,
        epochs=5,  # Small number for testing
        verbose='progress_bar'
    )
    model.fit(polymers)
    vectors = model.encode(polymers)
    print(f"Representation shape: {vectors.shape}")
    print(f"Representation for new molecule: {vectors[0]}")

if __name__ == "__main__":
    test_moama_encoder_polymers()
    test_moama_encoder()

Methods¶

`fit`	Fit the model to the training data with optional validation set.
`encode`	Encode molecules into vector representations.

For shared checkpoint and parameter operations, see Base Modules.

fit¶

MoamaMolecularEncoder.fit(X_train)[source]¶

Fit the model to the training data with optional validation set.

Parameters:: X_train (List[str]) – Training set input molecular structures as SMILES strings
Returns:: self – Fitted estimator
Return type:: MoamaMolecularEncoder

encode¶

MoamaMolecularEncoder.encode(X, return_type='pt')[source]¶

Encode molecules into vector representations.

Parameters:

X (List[str]) – List of SMILES strings
return_type (Literal["np", "pt"], default="pt") – Return type of the representations

Returns:

representations – Molecular representations

Return type:

ndarray or torch.Tensor