AutoEncoders for MD Trajectory Analysis

What this workflow does

This workflow trains a multilayer AutoEncoder on molecular dynamics trajectory data. The AutoEncoder serves three purposes: feature extraction, dynamical differences mapping, and collective variable generation for enhanced sampling.

The example system is the SARS-CoV-2 replicase polyprotein 1ab (PDB: 6W9C). The workflow fetches an apo trajectory and a holo trajectory from the MDDB database. It fits and featurizes both as Cα Cartesian coordinates.

The workflow builds an AutoEncoder architecture and trains it on the apo dataset. It then evaluates the trained model on both datasets. This comparison shows how the apo and holo latent spaces differ.

The workflow reconstructs the holo trajectory from the latent space. It computes RMSF profiles for the apo, holo, and reconstructed-holo trajectories. The three profiles tell you the reconstruction quality.

The last stage exports the trained model as a PLUMED collective variable, with the matching bias input files.

The compute problem

This workflow has the sharpest hardware split in the collection. One stage wants a GPU. Everything else wants a CPU.

The train_model stage runs PyTorch backpropagation over thousands of frames for many epochs. On a CUDA GPU it takes minutes. On a laptop CPU it takes hours. The epoch count and the latent dimension both push that number up.

The fetch_train_trajectory and fetch_test_trajectory stages are network-bound. They pull whole trajectories from MDDB. Bandwidth sets their run time, not the CPU.

The fitting and featurizing stages stream a trajectory through GROMACS and a featurizer. They are I/O-bound and single-threaded.

The RMSF stages and the PLUMED export are seconds of CPU work.

Run all of this on a GPU box and you rent GPU time for a large download and for several minutes of trajectory I/O. Run all of it on a laptop and you wait hours for the one stage that needs a GPU.

How Horus solves it

Horus assigns an executor per stage. The fetch, fit, featurize, RMSF, and PLUMED stages run on your local CPU. The train_model stage runs on a remote GPU box. Horus moves the featurized dataset across the boundary and brings the trained model back. You do not write copy commands.

The build_model stage is worth noting. It only writes the untrained architecture. It stays local. Only train_model and evaluate_model benefit from the GPU.

To change GPU providers, edit the executor: field on train_model. The runtime.command string is identical for a local CPU run, a remote CUDA node, and an HPC scheduler. You tune the architecture in configs/build_model.yaml on a laptop, then move the trained run to the GPU without touching the workflow logic.

biobb_pytorch, biobb_gromacs, and biobb_analysis all have native osx-arm64 conda builds. The full workflow runs on Apple Silicon in the conda executor. Only the training speed argues for a remote GPU.

Pipeline

Training (apo, 6W9C)
fetch_train_trajectory       Download apo trajectory 6W9C_apo from MDDB (mddb)
   │
fit_train_trajectory          Fit/align trajectory, rot+trans (gmx_image)
   │
featurize_train_trajectory     Featurize as Cα Cartesian coordinates (mdfeaturizer)
   │
build_model ──► train_model     Build AutoEncoder architecture, train on apo dataset

Evaluation (holo, 6W9C)
fetch_test_trajectory        Download holo trajectory 6W9C_holo from MDDB (mddb)
   │
fit_test_trajectory           Fit/align to apo reference, rot+trans (gmx_image)
   │
featurize_test_trajectory      Featurize as Cα Cartesian coordinates (mdfeaturizer)
   │
evaluate_model                  Evaluate trained AE on apo dataset (evaluate_model)
   │
reconstruct_trajectory           Reconstruct holo trajectory from AE latent space (feat2traj)

RMSF comparison
make_ndx_apo / make_ndx_holo     GROMACS Cα index files (make_ndx)
gmx_rmsf_apo / gmx_rmsf_holo / gmx_rmsf_holo_recon   RMSF of apo, holo, and reconstructed-holo

Enhanced sampling export
make_plumed                      Convert trained model to PLUMED CV (.ptc) + bias input files

Inputs and outputs

Inputs

The workflow has no file inputs. It fetches the apo and holo 6W9C trajectories from MDDB. configs/fetch_train_trajectory.yaml and configs/fetch_test_trajectory.yaml set the project and node IDs, the trajectory format, the frame range, and the atom selection.

Outputs land in results/:

• 6W9C_apo.pdb, 6W9C_apo.xtc, 6W9C_apo_fit.xtc: the apo training trajectory, raw and fitted.
• 6W9C_apo.pt, 6W9C_apo_stats.pt: the featurized apo dataset and its statistics.
• model.pth: the untrained AutoEncoder architecture.
• apo_trained_model.pth, model_training_metrics.npz: the trained model and the training curves.
• 6W9C_holo.pdb, 6W9C_holo.xtc, 6W9C_holo_fit.xtc: the holo test trajectory, raw and fitted.
• 6W9C_holo.pt, 6W9C_holo_stats.pt: the featurized holo dataset.
• apo_eval_results.npz: the evaluation and latent-space results.
• 6W9C_apo.ndx, 6W9C_holo.ndx: the GROMACS Cα index files.
• 6W9C_apo_rmsf.xvg, 6W9C_holo_rmsf.xvg, 6W9C_holo_recon_rmsf.xvg: the three RMSF profiles.
• 6W9C_holo_recon.xtc, 6W9C_holo_recon.pdb: the reconstructed holo trajectory.
• plumed_model.ptc, plumed.dat, features.dat: the PLUMED collective variable and its input files.

Run the workflow

Install the horus-runtime one time:

uv sync

If you do not have uv, install it first:

curl -LsSf https://astral.sh/uv/install.sh | sh

You can also install the packages with pip:

pip install horus-runtime horus-environments

Then run the workflow:

uv run horus run workflow.yaml

The first run builds the conda environment with PyTorch and GROMACS.

References

• BioExcel Building Blocks (biobb)
• biobb_pytorch
• PLUMED
• MDDB: MoDEL Database
• BioExcel AutoEncoder tutorial