Skip to content

jdsanc/ViSMA

Repository files navigation

ViSMA

Vibrational Spectroscopy Molecular Alignment

A framework for scaling vibrational spectroscopy respresenation learning for downstream tasks

PyTorch Lightning Hydra


Features

  • Flexible instantiation of architecture for vibrational spectra
  • Self-supervised tasks
  • Multi-GPU training support
  • Configurable preprocessing and augmentations

Installation

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 .

Quick Start

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()

Training

# 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.ckpt

Experiment configs live in configs/experiment/. See the Hydra docs for more on config composition.

Data

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.

Citation

@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}
}

License

MIT


Jesus Diaz Sanchez · jdsan@mit.edu · Learning Matter Lab, MIT

About

Vibrational Spectroscopy Molecular Alignment

Topics

Resources

Stars

Watchers

Forks

Releases

No releases published

Packages

 
 
 

Contributors