gmm2des

Fitting 2D Electronic Spectroscopy (2DES) via Spectral Density Prediction

gmm2des is a PyTorch-based framework for fitting simulated or experimental data by optimizing a Gaussian mixture model (GMM) to predict the system's spectral density function $J(\omega)$. The spectral density encodes how the bath couples to the electronic degrees of freedom and, via Mukamel's second-order cumulant expansion which is employed here, fully determines the linear and 2D electronic spectra.

---

Table of Contents

• Theoretical Background
• Architecture Overview
• Gaussian Mixture Model
• Loss Functions
• Configuration System
• Installation
• Data & Git LFS
• Usage
• Examples
• Units and Conventions
• Project Structure

---

Theoretical Background

2D Electronic Spectroscopy measures third-order nonlinear optical responses as a function of three time intervals $(t_1,t_2,t_3)$. The 2DES signal $S(t_1,t_2,\omega_3)$ is related to microscopic system-bath coupling through the spectral density $J(\omega)$, which describes the frequency distribution of bath modes weighted by their coupling strength.

Within Mukamel's second-order cumulant (Gaussian fluctuation) approximation, the lineshape function $g(t)$ is obtained from $J(\omega)$ by numerical integration of:

$$ g(t)=\frac{1}{\pi}\int_0^\infty\mathrm{d}\omega\frac{J(\omega)}{\omega^2}\left\lbrace\coth\left(\frac{\beta\hbar\omega}{2}\right)\left[1-\cos(\omega t)\right] - i\left[\omega t-\sin(\omega t)\right]\right\rbrace $$

The four Liouville-pathway response functions $(R_1,R_2,R_3,R_4)$ are then constructed from $g(t)$ evaluated at combinations of the three time intervals, and the 2DES $S(t_1,t_2,\omega_3)$ spectrum is obtained by Fourier transformation along $t_3$. The frequency-frequency correlation map $S(\omega_1,t_2,\omega_3)$ is obtained by a second Fourier transform across $t_1$.

This framework addresses the inverse problem: given an observed 2DES dataset, we aim to optimize the parameters of $J(\omega)$ such that the resulting theoretical spectrum best matches the data.

---

Architecture Overview

Trainable Model Parameters
        │
        ▼
Spectral Density J(ω)
        │
        ▼
Lineshape Function g(t)
        │
        ├──────────────────────┐
        ▼                      ▼
  Response Functions     Linear Absorption Spectrum
  R₁, R₂, R₃, R₄
        │
        ▼ (FFT over t₁ and t₃)
  2DES Spectrum S(t₁, t₂, ω₃)
        │
        ▼
  Loss vs. Reference Data
        │
        ▼
  Gradient Update

---

GMM - Gaussian Mixture Model

The GMM represents $J(\omega)$ as a weighted sum of Gaussians:

$$ J(\omega)=\sum_i a_i\exp\left[-\frac{1}{2}\left(\frac{\omega-\mu_i}{\sigma_i}\right)^2\right] $$

Trainable parameters: centers $\mu_i$, spreads $\sigma_i$, amplitudes $a_i$, and wegav $\bar{\omega}_{eg}$ (average electronic energy gap).

---

Loss Functions

Defined in src/gmm2des/training/loss_fns.py. All are composable via the TrainingLossFunction wrapper, which maps metric names to instantiated loss functions and applies optional per-metric weights.

Class	Description
SSIM	Structural Similarity Index
RootMeanSquaredError	RMSE

Loss configurations live in configs/training/loss/, and the default is ssim.yaml. Use ssim_with_linear.yaml to have the linear absorption spectrum contribute to the loss.

---

Configuration System

The project uses Hydra for hierarchical YAML configuration. The root config is src/gmm2des/configs/train.yaml, which composes the following config groups:

train.yaml
├── data/train:        multidim_default
├── model:             gmm
├── training/
│   ├── checkpointing: default
│   ├── dataloader_args: default
│   ├── logging:       default (TensorBoard) | wandb
│   ├── loss:          ssim | ssim_with_linear
│   └── training_args: default
└── global_args:
    ├── device: cuda
    └── seed: null

Override any group or value on the command line using standard Hydra syntax. For example:

# Use W&B logging, 2DES + lin. abs. loss, and a custom data file
gmm2des_train training/logging=wandb training/loss=ssim_with_linear data.train.filename=/path/to/data.h5

---

Installation

Requirements: Python ≥ 3.11, PyTorch ≥ 2.0.

# Install from source
pip install -e .

# With uv (recommended for development)
uv sync

Development dependencies (linting, type checking, tests):

uv sync --group dev

The package installs two CLI entry points: gmm2des_train and gmm2des_loss_scan.

---

Data & Git LFS

The 2DES reference datasets (data/*/data2d/*.h5, ~2.3 GB total) are tracked with Git LFS (see .gitattributes) rather than committed directly, since they're too large for a normal git object.

Install Git LFS once per machine (before cloning, if possible):

git lfs install

If you already cloned before installing LFS, the .h5 files will be small text pointers instead of real data — fetch them with:

git lfs pull

To skip downloading the .h5 files (e.g. you only need the configs, code, or pretrained checkpoints in data/trained_models/), clone without smudging LFS pointers, then pull only what you need:

GIT_LFS_SKIP_SMUDGE=1 git clone <repo-url>
cd gmm2des

# Pull everything later:
git lfs pull

# ...or pull just one dataset:
git lfs pull --include="data/gfp_water/data2d/*.h5"

Training and the loss scan script both expect the actual .h5 file to be present at data.train.filename; running against an un-smudged pointer file will fail to load as HDF5.

---

Usage

Training

# Train with default settings (GMM model, SSIM loss, CUDA), using one of the bundled datasets
gmm2des_train data.train.filename=data/gfp_water/data2d/gfp_water.h5

# Train with a specific loss and number of Gaussians
gmm2des_train training/loss=ssim_with_linear model.num_gaussians=50 data.train.filename=data/gfp_water/data2d/gfp_water.h5

# Train on CPU
gmm2des_train global_args.device=cpu data.train.filename=data/gfp_water/data2d/gfp_water.h5

# Enable profiling (outputs TensorBoard trace to logs/trace/)
PROFILE_RUN=1 gmm2des_train data.train.filename=data/gfp_water/data2d/gfp_water.h5

Each run creates a versioned subdirectory under logs/version_N/, containing the full resolved config (full_config.yaml) and model checkpoints.

For launching training on a SLURM cluster, examples/submit_scripts/ contains one example submission script per bundled dataset (gfp_water.sh, nileblue_ethanol.sh, nilered_benzene.sh, pyp_vacuum.sh, experiment_no_lin.sh, experiment_with_lin.sh). Each script needs a scheduler header added (YOUR HEADER GOES HERE) and the data_dir variable pointed at your local copy of data/ before submitting.

Programmatic Use

from gmm2des import GMM, buildModel, loadSystemData
from omegaconf import OmegaConf

# Load a 2DES dataset
data_cfg = OmegaConf.create({"filename": "data/gfp_water/data2d/gfp_water.h5", "t2_values": None})
dataset = loadSystemData(data_cfg)

# Build a GMM model
model_cfg = OmegaConf.load("src/gmm2des/configs/model/gmm.yaml")
model = buildModel(model_cfg, reference_data=dataset, device="cuda")

# Freeze the lineshape and inject a custom spectral density
model.freeze()
model.spectral_density = my_custom_jw
model.lineshape = my_precomputed_gt

See examples/parse_results.py and examples/results.ipynb for a complete, ready-to-run example of loading a trained model and generating figures (see Examples below).

Loss Scan Script

# Sweep over population times defined in loss_scan.yaml; override with standard Hydra syntax
gmm2des_loss_scan

---

Examples

The examples/ directory contains ready-to-run scripts for visualizing trained models, on top of the pretrained checkpoints shipped in data/trained_models/.

File	Purpose
examples/parse_results.py	Loads one or more trained model versions from data/trained_models/<system>/version_N and produces the spectral density (with linear absorption inset), 2DES spectra at requested population times, and intensity trace figures. Configured via Hydra using plot_cfg.yaml; can be run as a script (python examples/parse_results.py) or imported and called via its main(cfg) function.
examples/results.ipynb	A thin notebook wrapper around parse_results.py. Set main_dir to your local clone of this repository and system_name to one of expt_no_lin, expt_with_lin, gfp, nilered, or pyp, then run all cells to reproduce the corresponding figures.
examples/plot_cfg.yaml	The Hydra config consumed by parse_results.py (device, which run(s) to load, plotting options for each figure type).
examples/submit_scripts/	Example SLURM submission scripts for training each of the bundled datasets from scratch (see Training).

For example, to view the results for a single version trained on GFP in water:

cd examples
python parse_results.py load_info.run_dirs=[../data/trained_models/gfp_water] load_info.filter_options=[{version:0}]

---

Units and Conventions

When specifying axes (e.g., spectral_density_axis) or training on 2DES data, the framework expects the following as inputs:

Quantity	Unit
Time $(t_1,t_2,t_3)$	femtoseconds (fs)
Frequency ($\omega$)	$\mathrm{cm}^{-1}$
Energy $(\bar{\omega}_{eg},\omega_3)$	electronvolts (eV)
Temperature	300 K

However, internally the framework operates by converting both the frequency units ($\textrm{cm}^{-1}$) and energy units (eV) to frequency units consistent with the time units (1/fs). Constants are initialized at runtime via consts.initConstants("eV,fs,cm,cm^-1"). Conversion factors (e.g., thz_to_eV, wn_to_Hz) are available in src/gmm2des/utils/spectra/constants.py.

---

Project Structure

src/gmm2des/
├── configs/                   # Hydra configuration files
│   ├── train.yaml             # Root config
│   ├── model/                 # Model configs
│   ├── data/train/            # Dataset configs
│   └── training/              # Loss, optimizer, logging, checkpointing configs
│
├── model/                     # Model definitions
│   ├── abstract_model.py      # AbstractSpectrumModel base class
│   ├── gmm.py                 # Gaussian Mixture Model
│
├── training/                  # Training infrastructure
│   ├── train.py               # runEpoch / train loop
│   ├── loss_fns.py            # All loss function implementations
│   ├── checkpointer.py        # Checkpoint saving/loading
│   └── logger.py              # TensorBoard / W&B logger wrapper
│
├── data/                      # Dataset classes
│   ├── abstract_datasets.py   # SpectrumDataset base class
│   ├── multidim.py            # MultidimSpectrum (2DES)
│   ├── linear.py              # LinearSpectrum + PulseProfile
│   └── system.py              # System-level dataset wrapper
│
├── utils/
│   ├── spectra/
│   │   ├── constants.py       # Physical constants + unit initialization
│   │   ├── fourier.py         # FFT / iFFT / ttwToWtw utilities
│   │   └── physical_calculations.py  # lineshapeFromSpectralDensity, responseFunctions, multidimSpectrum
│   ├── model/
│   │   ├── build_model.py     # Model factory from config
│   │   ├── forward_fns.py     # realMultidimSpectrum, realMultidimSpectrumWithLinearAbsorption, etc.
│   │   ├── spectral_density_fns.py   # GMM → J(ω) functions
│   │   ├── init_model_params.py      # Parameter initialization from config
│   │   └── model_map.py       # @defineModel registry decorator
│   ├── data/
│   │   ├── load_data.py       # loadSystemData / loadLinearData from HDF5
│   │   ├── split_data.py      # Train/val/test splitting
│   │   ├── forward_fns.py     # Data preprocessing forward functions
│   │   ├── interpolation.py   # Spectral interpolation utilities
│   │   └── spectra_utils.py   # Normalization helpers
│   └── training/
│       ├── parse_cfg.py       # Parse Hydra config into train() arguments
│       └── parse_train_fwd_args.py  # Parse forward function arguments
│
└── scripts/
    ├── run_train.py           # gmm2des_train CLI entry point
    └── loss_scan.py           # Hyperparameter sweep script

data/                           # Data used in the main text
├── gfp_water/
├── nileblue_ethanol/
├── nilered_benzene/
├── pyp_vacuum/
├── experiment/
├── trained_models/             # Pretrained checkpoints + full_config.yaml, one subdirectory per system above
├── ir_prob_dist.npy
└── loss_scan_times.npy

examples/                       # Tracked quickstart scripts
├── parse_results.py            # Load a trained model and generate figures
├── results.ipynb               # Notebook wrapper around parse_results.py
├── plot_cfg.yaml               # Hydra config used by parse_results.py
└── submit_scripts/             # Example SLURM submission scripts, one per bundled dataset

Each of gfp_water/, nileblue_ethanol/, nilered_benzene/, pyp_vacuum/, and experiment/ shares the same layout: data1d/ holds the reference linear_absorption.pkl and spec_dens_eVfs.pkl (plus pulse_spectral_profile.pkl for experiment/), and data2d/ holds the <system>.h5 2DES dataset used for training.