BioEmu-CV

This is the official implementation of "Learning Collective Variables from BioEmu with Time-Lagged Generation" (Park et al., ICLR 2026). [arXiv]

A research variant of BioEmu tailored for learning machine learning collective variables (MLCVs) from protein structure ensembles. Built on top of the BioEmu diffusion model, it integrates with mlcolvar to train low-dimensional collective variables that can be used to drive enhanced sampling simulations.

Table of Contents

• Environment Setup
• Installation
• Training a CV Model
• Running OPES Simulations
• Analysis

---

Environment Setup

OPES simulations require GROMACS patched with PLUMED. The recommended way to set up this environment is via the Docker container provided at:

https://github.com/seonghyun26/gromacs-plumed-docker

This container ships GROMACS pre-patched with PLUMED and supports NVIDIA GPUs (requires CUDA 11.2 and driver ≥ 450.80.02).

Pull and run:

# Run a GROMACS command with GPU support
docker run --gpus all -u $(id -u) -w /work -v $PWD:/work \
    <image>:<tag> gmx mdrun ...

Note: Make sure gmx resolves to the PLUMED-patched binary inside the container before running OPES simulations. The opes/main.py script calls gmx mdrun directly, so the container environment (or a matching host installation) must be active when executing it.

---

Installation

Install the package and its dependencies from the repository root:

pip install -e .

Additional dependencies used by this research variant:

pip install mlcolvar pytorch-lightning wandb hydra-core

Note: On first use, BioEmu will download ColabFold for MSA/embedding generation into ~/.bioemu_colabfold. Set BIOEMU_COLABFOLD_DIR to override the location.

---

Training a CV Model

CV models are trained from the repository root using control.py with Hydra configs located in config/.

Supported molecules: cln025, 2jof, 1fme

Basic training command

CUDA_VISIBLE_DEVICES=0 python control.py --config-name score-cln025

Each score-<molecule>.yaml config composes from sub-configs for data, model, logging, and sampling:

config/
├── score-cln025.yaml      # top-level config (data + model + log + sample)
├── score-2jof.yaml
├── score-1fme.yaml
├── ...

Using the provided scripts

cd scripts/
bash ctrl-score-cln025.sh <GPU_ID>   # e.g., bash ctrl-score-cln025.sh 0

Overriding config values

CUDA_VISIBLE_DEVICES=0 python control.py \
    --config-name score-cln025 \
    ++log.date=$(date +%m%d_%H%M%S) \
    ++log.tags="['pilot','score']"

Output

Trained models are saved as TorchScript JIT files under model/:

model/
└── <datetime>-<MOLECULE>-jit.pt

These .pt files are loaded directly by PLUMED during OPES simulations.

---

Running OPES Simulations

OPES simulations are managed by opes/main.py. Configs and PLUMED .dat files live in opes/config/, and per-molecule MD input files (.tpr, .top, etc.) live in opes/data/<MOLECULE>/.

Available CV methods

Method	Description
ours	BioEmu-trained MLCV
tica	Time-lagged Independent Component Analysis
tda	Time-lagged Deep Autoencoder
tae	Time-lagged Autoencoder
vde	Variational Dynamics Encoder

Running a simulation

cd opes/
python main.py --config-name ours-cln025

Config files follow the pattern <method>-<molecule>.yaml, e.g.:

opes/config/
├── ours-cln025.yaml
├── tica-2jof.yaml
├── tda-1fme.yaml
└── ...

Key config options (ours-cln025.yaml)

molecule: cln025
method: ours
ckpt_dir: '.'              # directory under opes/model/ containing the JIT model
ckpt_path: <datetime>-CLN025  # JIT model filename (without -jit.pt)
start_gpu: 0               # first GPU index
opes:
  step: 500_000_000        # total MD steps
  sigma: 0.05              # OPES kernel bandwidth
  max_seed: 3              # number of parallel replicas (seeds 0..max_seed)
  ntomp: 1                 # OMP threads per replica

Parallel execution

The runner launches max_seed + 1 independent GROMACS processes in parallel (one per GPU), then waits for all to complete. Each replica writes outputs to:

opes/simulations/<molecule>/<method>/<datetime>/<seed>/

Simulation progress and PLUMED parameters are logged to Weights & Biases under the opes project.

Note: Ensure gmx is available in PATH (via the Docker container or a host GROMACS+PLUMED installation) before running simulations.

---

Analysis

OPES trajectory analysis

cd opes/
bash scripts/run_analysis.sh <method> <molecule> <datetime>

# Example
bash scripts/run_analysis.sh tda cln025 0818_150938

Or directly:

python analysis_opes.py --config-name tda-cln025 \
    date=<datetime> \
    analysis.gmx=True

CV projection analysis (from root)

python analysis_cv.py

Results and plots are logged to W&B and saved locally under img/.