Vibrational Spectroscopy Molecular Alignment
A framework for scaling vibrational spectroscopy respresenation learning for downstream tasks
- Flexible instantiation of architecture for vibrational spectra
- Self-supervised tasks
- Multi-GPU training support
- Configurable preprocessing and augmentations
git clone https://github.com/jdsanc/ViSMA && cd ViSMA
conda create -n spectra python=3.10 && conda activate spectra
pip install torch torchvision --index-url https://download.pytorch.org/whl/cu118
pip install -r requirements.txt && pip install -e .from src.models.visma_module import ViSMALitModule
model = ViSMALitModule.load_from_checkpoint("path/to/checkpoint.ckpt")
model.eval()
# Get embeddings
spec_embedding = model.encode_spectra(spectra) # [batch, 2, 3501] → [batch, 256]
mol_embedding = model.encode_molecule("CCO") # SMILES → [batch, 256]
# Compute similarity
similarity = (spec_embedding @ mol_embedding.T).squeeze()# Train on CPU
python src/train.py trainer=cpu
# Train on GPU
python src/train.py trainer=gpu
# Multi-GPU
python src/train.py trainer=gpu trainer.devices=4 trainer.strategy=ddp
# Run a specific experiment
python src/train.py experiment=experiment_name
# Override parameters
python src/train.py experiment=ViSMA \
model.optimizer.lr=5e-5 \
data.batch_size=64
# Resume from checkpoint
python src/train.py experiment=ViSMA \
ckpt_path=logs/runs/.../checkpoints/last.ckptExperiment configs live in configs/experiment/. See the Hydra docs for more on config composition.
Spectra should be numpy arrays with shape [N, channels, wavenumbers]. You can configure which channels (IR, Raman, or both) and the wavenumber range at model instantiation.
import numpy as np
spectra = np.stack([ir, raman], axis=1) # [N, 2, 3501]
np.save('data/dataset_name/spectra.npy', spectra)Preprocessing and augmentation options (line broadening, normalization, etc.) are configurable via dataclass configs. See USERGUIDE.md for details.
@software{diaz_sanchez_2026_visma,
author = {Diaz Sanchez, Jesus},
title = {{ViSMA}: Vibrational Spectroscopy Molecular Alignment},
year = {2026},
url = {https://github.com/jdsanc/ViSMA},
version = {0.1.0}
}MIT
Jesus Diaz Sanchez · jdsan@mit.edu · Learning Matter Lab, MIT