SANDMolecularEncoder

Pairwise similarity scores; dense or compressed molecular descriptors.

from torch_molecule import SANDMolecularEncoder

Methods

Constructor

class SANDMolecularEncoder(*, num_layer=7, hidden_size=1024, emb_dim=512, encoder_type='gine', num_atom_features=62, num_edge_features=16, normalize_shared_emb=True, quantize_emb=True, num_ivf_quantizers=1, num_pq_quantizers=32, ivf_codebook_size=256, ivf_threshold_ema_dead_code=0.5, ivf_ema_decay=0.8, similarity_loss_weight=1.0, regularization_strength=1.0, uniformity_loss_weight=0.5, uniformity_margin=0.2, uniformity_num_matches=1, ivf_loss_weight=0.5, pq_loss_weight=0.25, batch_size=128, epochs=100, learning_rate=0.0005, min_lr=1e-05, weight_decay=0.05, warmup_iters=10000, decay_iters=1000000, grad_clip_value=None, emb_key=None, device=None, model_name='SANDMolecularEncoder', verbose='none')

SAND: Shape-Aware Neural Descriptor.

Parameters:

• num_layer (int, default=7) – Number of message-passing layers.

• hidden_size (int, default=1024) – GNN hidden width.

• emb_dim (int, default=512) – Molecular embedding dimension.

• encoder_type (str, default="gine") – GNN architecture: gine, gatv2, or gps.

• num_atom_features (int, default=62) – Atom feature vocabulary size; at least 62 for the fixed RDKit features.

• num_edge_features (int, default=16) – Edge feature vocabulary size; at least 16 for the fixed RDKit features.

• normalize_shared_emb (bool, default=True) – Normalize backbone embeddings during training and quantized reconstructions.

• quantize_emb (bool, default=True) – Enable IVF/PQ compression and both quantization losses. False selects dense training.

• num_ivf_quantizers (int, default=1) – Number of IVF quantizer heads.

• num_pq_quantizers (int, default=32) – Number of product-quantization subspaces.

• ivf_codebook_size (int, default=256) – Number of IVF centroids per head.

• ivf_threshold_ema_dead_code (float, default=0.5) – EMA count threshold for replacing unused IVF codes.

• ivf_ema_decay (float, default=0.8) – IVF centroid exponential-moving-average decay.

• similarity_loss_weight (float, default=1.0) – Weight of the soft-Spearman similarity loss.

• regularization_strength (float, default=1.0) – Temperature of the differentiable rank function.

• uniformity_loss_weight (float, default=0.5) – Weight of the embedding uniformity loss.

• uniformity_margin (float, default=0.2) – Cosine-similarity margin for the uniformity penalty.

• uniformity_num_matches (int, default=1) – Number of matches per query; must be 1.

• ivf_loss_weight (float, default=0.5) – Weight of the IVF cosine-alignment loss.

• pq_loss_weight (float, default=0.25) – Weight of the PQ squared-distance loss.

• batch_size (int, default=128) – Target pairs per training batch (at least 2), or molecules per encoding batch.

• epochs (int, default=100) – Number of training epochs.

• learning_rate (float, default=0.0005) – Peak learning rate.

• min_lr (float, default=0.00001) – Minimum rate after cosine decay.

• warmup_iters (int, default=10000) – Number of optimizer steps for learning-rate warmup.

• decay_iters (int, default=1000000) – Optimizer-step horizon for warmup and cosine decay.

• weight_decay (float, default=0.05) – AdamW weight decay.

• grad_clip_value (float, optional) – Maximum gradient norm. None disables clipping.

• emb_key (str, optional) – Embedding to return: emb or emb_quantized. None selects automatically.

• device (torch.device or str, optional) – CPU or CUDA device; defaults to CUDA when available.

• model_name (str, default="SANDMolecularEncoder") – Estimator name.

• verbose (str, default="none") – One of none, progress_bar, or print_statement.

References

Paper: https://proceedings.mlr.press/v306/winter26a.html Code: https://github.com/pfizer-opensource/SAND

Methods

fit	Fit on molecular-pair similarity scores.
encode	Return molecular embeddings as CPU tensors or NumPy arrays.
load_from_pretrained	Load pretrained SAND weights.
save_to_local	Save model to local disk.
load_from_local	Load a checkpoint, preserving the current encoder if loading fails.

For shared checkpoint and parameter operations, see Base Modules.

fit

SANDMolecularEncoder.fit(X_train, y_train=None, *, pairs=None)

Fit on molecular-pair similarity scores.

Parameters:

• X_train (list of str) – Training SMILES strings.

• y_train (array-like of shape (n_pairs,)) – Similarity score for each pair.

• pairs (array-like of shape (n_pairs, 2)) – Query and match indices into X_train.

Returns:

self – Fitted encoder.

Return type:

SANDMolecularEncoder

encode

SANDMolecularEncoder.encode(X, return_type='pt')

Return molecular embeddings as CPU tensors or NumPy arrays.

Parameters:

• X (List[str])

• return_type (Literal['np', 'pt'])

Return type:

ndarray | Tensor

load_from_pretrained

SANDMolecularEncoder.load_from_pretrained(checkpoint_path)

Load pretrained SAND weights.

Parameters:

checkpoint_path (str) – Checkpoint file matching the model architecture.

Returns:

self – Encoder with loaded weights.

Return type:

SANDMolecularEncoder

save_to_local

SANDMolecularEncoder.save_to_local(path)

Save model to local disk.

Parameters:

path (str) – File path to save the model

Raises:

ValueError – If the model is not fitted

Return type:

None

load_from_local

SANDMolecularEncoder.load_from_local(path)

Load a checkpoint, preserving the current encoder if loading fails.

Parameters:

path (str)

Return type:

None