MassSpecGym

MassSpecGym provides a modular infrastructure for benchmarking molecular discovery tasks. The workflow typically follows this pattern:

1. Data Loading: Use load_massspecgym() to get the dataset as a pandas DataFrame, or use MassSpecDataset to load it directly from Hugging Face into a PyTorch-compatible format.
2. Preprocessing: Apply transforms (like SpecTokenizer for spectra or MolFingerprinter for molecules) via the MassSpecDataset class.
3. Data Management: Use MassSpecDataModule (a PyTorch Lightning LightningDataModule) to handle data splitting (train/val/test) and batching.
4. Model Implementation: Inherit from a challenge-specific abstract class (e.g., RetrievalMassSpecGymModel) and implement the forward and step methods to define your architecture and training/inference logic.
5. Training/Evaluation: Use a standard PyTorch Lightning Trainer to run the training and testing loops.

from massspecgym.data import MassSpecDataset
from massspecgym.data.transforms import SpecTokenizer, MolFingerprinter
from massspecgym.data import MassSpecDataModule

# 1. Define transforms and dataset
dataset = MassSpecDataset(
    spec_transform=SpecTokenizer(n_peaks=60),
    mol_transform=MolFingerprinter(),
)

# 2. Initialize data module for Lightning
data_module = MassSpecDataModule(
    dataset=dataset,
    batch_size=32
)

To implement a custom model, you must inherit from one of the three challenge-specific abstract classes:

• DeNovoMassSpecGymModel (for de novo molecule generation)
• RetrievalMassSpecGymModel (for molecule retrieval)
• SimulationMassSpecGymModel (for spectrum simulation)

You are required to implement:

1. forward(self, x): Defines the prediction logic (e.g., mapping spectra to a fingerprint).
2. step(self, batch, stage): Defines the custom logic for training and inference, including unpacking the batch, calculating loss, and computing similarity scores.

from massspecgym.models.retrieval.base import RetrievalMassSpecGymModel
from massspecgym.models.base import Stage
import torch
import torch.nn as nn

class MyModel(RetrievalMassSpecGymModel):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        # Define architecture here

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # Implement prediction logic
        return x

    def step(self, batch: dict, stage: Stage) -> tuple[torch.Tensor, torch.Tensor]:
        # Unpack inputs from batch
        x = batch["spec"]
        fp_true = batch["mol"]
        cands = batch["candidates"]
        batch_ptr = batch["batch_ptr"]

        # Predict
        fp_pred = self.forward(x)

        # Calculate loss and scores
        loss = nn.functional.mse_loss(fp_true, fp_pred)
        fp_pred_repeated = fp_pred.repeat_interleave(batch_ptr, dim=0)
        scores = nn.functional.cosine_similarity(fp_pred_repeated, cands)

        return dict(loss=loss, scores=scores)

To quantify data leakage, calculate the maximum Tanimoto similarity between molecules in the test set and all molecules in the training/validation sets. High Tanimoto similarity (e.g., > 0.95) between test and training molecules indicates potential leakage where the model might be memorizing similar structures rather than learning generalizable features.

This process typically involves:

1. Generating Morgan fingerprints for all molecules using massspecgym.utils.morgan_fp.
2. Using rdkit.Chem.DataStructs.FingerprintSimilarity to compute similarities.
3. Iterating through the test set to find the maximum similarity score relative to the training/validation pool.

from rdkit import Chem
from rdkit.Chem import DataStructs
import massspecgym.utils as utils

# Logic for finding max similarity
def max_train_tanimoto(test_mol_and_fp, train_mols_and_fps):
    tanimotos = [
        DataStructs.FingerprintSimilarity(test_mol_and_fp[1], train_mol_and_fp[1])
        for train_mol_and_fp in train_mols_and_fps
    ]
    return max(tanimotos)

To clean a dataset of duplicate mass spectra, you can group spectra by their inchikey metadata and use similarity scoring to identify and remove redundant entries.

1. Group spectra by inchikey.
2. For each group, calculate similarity scores (e.g., using CosineGreedy from matchms).
3. Identify duplicates where the similarity score is exactly 1.0 (excluding the diagonal/self-comparison).
4. Retain only one unique spectrum per group.

from matchms.similarity.CosineGreedy import CosineGreedy
from matchms import calculate_scores
import numpy as np

# Assuming cleaned_spectra is a list of matchms Spectrum objects
# and unique_inchikeys is a dict mapping inchikey -> list of indices

for inchikey in unique_inchikeys:
    matching_spectra = [cleaned_spectra[idx] for idx in unique_inchikeys[inchikey]]
    
    # Calculate similarity scores within the group
    scores = calculate_scores(references=matching_spectra,
                              queries=matching_spectra,
                              similarity_function=CosineGreedy())
    
    scores_array = scores.to_array(name="CosineGreedy_score")
    
    # Find indices where similarity is 1.0 (duplicates)
    indices = np.argwhere(scores_array == 1.0)

    # Filter out diagonal indices (where index[0] == index[1]) to avoid self-matching
    non_diagonal_indices = [tuple(index) for index in indices if index[0] != index[1]]
    
    # Logic follows to remove identified duplicates from the list

The data pipeline is constructed using three main components:

1. RetrievalDataset: The core dataset object. It requires two transforms:
   • spec_transform: A transformation for the spectra (e.g., SpecTokenizer(n_peaks=...)).
   • mol_transform: A transformation for the molecules (e.g., MolFingerprinter(fp_size=...)).
2. MassSpecDataModule: Wraps the dataset for PyTorch Lightning. It takes the dataset, batch_size, and num_workers as arguments.
3. Transforms:
   • SpecTokenizer: Converts raw spectra into a tokenized format with a fixed number of peaks.
   • MolFingerprinter: Converts molecules into fixed-size fingerprints.

from massspecgym.data import RetrievalDataset, MassSpecDataModule
from massspecgym.data.transforms import SpecTokenizer, MolFingerprinter

n_peaks = 60
fp_size = 4096
batch_size = 32

# 1. Create the dataset with transforms
dataset = RetrievalDataset(
    spec_transform=SpecTokenizer(n_peaks=n_peaks),
    mol_transform=MolFingerprinter(fp_size=fp_size),
)

# 2. Create the data module
data_module = MassSpecDataModule(
    dataset=dataset,
    batch_size=batch_size,
    num_workers=4
)