This is an old revision of the document!
GMVAE Clustering of Protein Contacts
(author: Drs. Min-Kang Hsieh and Mahdi Ghorbani)
This is based off Dr. Ghorbani's GMVAE code available on Github. This has been updated to work with newer versions of python and Zaratan's architecture. Use on other HPC will likely need tweaking from the instructions below.
Tar file for all code: gmcave_tf1_update_copy.tgz
The detailed instruction was added by Dr. Min-Kang Hsieh.
This instruction includes a complete workflow for unsupervised clustering of protein conformations using pairwise Cα (alpha-carbon) distance matrices (or BB atoms in CG modeling) as the input representation for a Gaussian Mixture Variational Autoencoder (GMVAE). The workflow covers dataset preparation, model hyperparameter optimization, systematic grid search over the number of Gaussian mixture components (K) and latent-space dimensions (z), and latent-space evaluation to identify the optimal (K, z) configuration for conformational clustering.
In addition to GMVAE training and evaluation, the package provides utilities for generating cluster-wise contact maps, representative conformations, and summary statistics for structural analysis and visualization.
Recommended Computing Workflow and Resources
The workflow can be executed on either GPU or CPU systems. The following hardware configurations are recommended:
- Dataset preparation: CPU
- Single-model training (for Hyperparameter optimization): One NVIDIA V100 GPU (recommended) or ≥64 CPU cores for small datasets
- Grid search: One NVIDIA V100 GPU for training; ≥64 CPU cores for post-training evaluation (CPU-only grid search is suitable only for small datasets)
- Model evaluation: ≥64 CPU cores (typically performed after training)
- Cluster visualization (contact map and structural interpretation): CPU
Python Environment
Note: The current GMVAE workflow requires two independent Python environments. Separating the environments avoids dependency conflicts between PyEMMA (used for trajectory processing) and TensorFlow 1.15.4 (used for GMVAE training). This is recommended because PyEMMA depends on several compiled C/C++/Fortran libraries (e.g., MDTraj, NumPy, and SciPy), which may conflict with the TensorFlow environment.
Environment 1: Dataset Preparation (PyEMMA) and Cluster Contact Map This environment is used for molecular dynamics trajectory processing and dataset generation.
Option 1: Create from the provided environment file (recommended)
conda env create -p ~/scratch/pyemma_env -f pyemma.yml
Option 2: Install manually
conda create -p ~/scratch/pyemma_env -c conda-forge python=3.9 pyemma mdtraj
Activate the environment:
conda activate ~/scratch/pyemma_env
Environment 2: GMVAE Train/Grid/Eval (TensorFlow 1.15.4) This environment is used for GMVAE training, hyperparameter optimization, grid search, and model evaluation. Create the environment from the provided environment file:
conda env create -p ~/scratch/tf1.15.4 -f tf1.15.4.yml conda activate ~/scratch/tf1.15.4
Package Layout Tree
gmvae_tf1_v2/
├─ pyemma.yml
├─ tf1.15.4.yml
├─ src/
│ ├─ gmvae_test/
│ │ ├─gmvae_pipeline_v2.py
│ │ ├─real_conv_GMVAE_graph_mse.py
│ │ ├─real_Conv_GMVAE_gumbel.py
│ │ ├─load_data.py
│ │ ├─logger.py
│ │ ├─base_graph.py
│ │ └─base_model.py
│ └─ utils/
│ ├─data_collection_v4.py
│ └─contact_heatmap_top_structures.py
└─ scripts/
├─run_datacoll.csh
├─run_datacoll_list.csh
├─run_train.csh
├─run_grid.csh
├─run_grid_manually.csh
├─run_grid_manually.scr
├─run_eval.csh
└─run_contact_heatmap_top_structures.csh
Usage note: copy the scripts folder as your project name to begin.
Input data preparation
Modify run_datacoll.csh or run_datacoll_list.csh to suit your input trajectory files. Detailed instructions for using this data collection script are provided below.
usage: data_collection_v4.py [-h] --pdb PDB
[--traj_files TRAJ_FILES [TRAJ_FILES ...]]
[--traj_dirs TRAJ_DIRS [TRAJ_DIRS ...]]
[--traj_prefix TRAJ_PREFIX]
[--traj_ext {dcd,trr,xtc}] [--start START]
[--end END] [--single_chain] [--dimer]
[--selected_atom SELECTED_ATOM]
[--indices1 INDICES1 INDICES1]
[--indices2 INDICES2 INDICES2] [--skip SKIP]
[--skip_first_frame SKIP_FIRST_FRAME]
[--output_prefix OUTPUT_PREFIX]
optional arguments:
-h, --help show this help message and exit
--pdb PDB Path to topology/structure file, usually .pdb or .gro.
--traj_files TRAJ_FILES [TRAJ_FILES ...]
Explicit list of trajectory files, e.g. traj1.trr
traj2.trr.
--traj_dirs TRAJ_DIRS [TRAJ_DIRS ...]
List of directories containing trajectory files.
--traj_prefix TRAJ_PREFIX
Prefix for trajectory files when using traj_dirs.
--traj_ext {dcd,trr,xtc}
Trajectory extension when using traj_dirs.
--start START Starting trajectory index.
--end END Ending trajectory index.
--single_chain Use the same atom indices for both chains.
--dimer Use separate atom indices for two chains.
--selected_atom SELECTED_ATOM
Selected atom name for analysis.
--indices1 INDICES1 INDICES1
Range for chain 1 MDTraj atom indices, e.g. 0 99.
Range for chain 2 MDTraj atom indices, e.g. 100 220.
--skip SKIP Frame skip interval.
--skip_first_frame SKIP_FIRST_FRAME
Number of initial frames to discard after loading all
trajectories.
--output_prefix OUTPUT_PREFIX
Prefix for output .npy files.
GMVAE Clustering Workflow
The GMVAE workflow consists of three major stages:
- Hyperparameter optimization (excluding K and z)
- Grid search over the number of clusters (K) and latent dimensions (z)
- Latent-space evaluation to identify the optimal (K, z) configuration
Quick start
Activate the TensorFlow environment before running the GMVAE pipeline.
conda activate ~/ scratch/tf1.15.4
The default training hyperparameters are: –temp 0.05 –dropout 0.3 –lr 0.001 –epochs 300 –batch-size 100 –dense-n 64 –filters 64 64 32 –kernel-sizes 3 3 3 –pool-sizes 2 2 1 –num-layers 3
The gmvae_pipeline_v2.py script provides a unified interface for GMVAE training, hyperparameter grid search, and evaluation of GMVAE for data clustering.
The program supports three operating modes:
- train – Train a single GMVAE model.
- grid – Perform a grid search over multiple values of K and z.
- eval – Evaluate a trained model and generate clustering results.
Training Mode (train)
Train a single GMVAE model using a specified set of hyperparameters.
Example: Single Training
python gmvae_pipeline_v2.py --mode train \ --pdb /home/mhsieh/scratch/bst2/orf7a-cov-2_bst2_short/aa/c2_r1_cg2aa/openmm/ca_dcd/ca.pdb \ --dataset-npy orf7a_bst2_short_data.npy \ --rmsd-npy orf7a_bst2_short_rmsd.npy \ --mode_chain dimer \ --indices1 0 99 \ --indices2 100 220 \ --K 6 --z-dim 8 \ --epochs 100 --batch-size 500 \ --output-dir outputs
Example: Training with Automatic Evaluation
python gmvae_pipeline_v2.py --mode train \ --pdb /home/mhsieh/scratch/bst2/orf7a-cov-2_bst2_short/aa/c2_r1_cg2aa/openmm/ca_dcd/ca.pdb \ --dataset-npy orf7a_bst2_short_data.npy \ --rmsd-npy orf7a_bst2_short_rmsd.npy \ --mode_chain dimer \ --indices1 0 99 \ --indices2 100 220 \ --K 6 --z-dim 4 \ --epochs 100 --batch-size 500 \ --output-dir outputs \ --evaluate_after_train
Grid Search Mode (grid)
Once a suitable set of training hyperparameters has been determined (excluding K and z), perform a systematic grid search to identify the optimal latent-space dimension and number of clusters.
**Example** python gmvae_pipeline_v2.py --mode grid \ --pdb /home/mhsieh/scratch/bst2/orf7a-cov-2_bst2_short/aa/c2_r1_cg2aa/openmm/ca_dcd/ca.pdb \ --dataset-npy orf7a_bst2_short_data.npy \ --rmsd-npy orf7a_bst2_short_rmsd.npy \ --mode_chain dimer \ --indices1 0 99 \ --indices2 100 220 \ --grid-ks 3 4 5 6 7 8 \ --grid-zs 5 6 7 \ --epochs 100 --batch-size 500 \ --output-dir grid_outputs
Evaluation Mode (eval)
Evaluation is performed after model training and is primarily CPU-based. The evaluation step generates latent-space statistics and clustering summaries for model selection.
Example
python gmvae_pipeline_V2.py --mode eval \ --pdb /home/mhsieh/scratch/bst2/orf7a-cov-2_bst2_short/aa/c2_r1_cg2aa/openmm/ca_dcd/ca.pdb \ --dataset-npy orf7a_bst2_short_data.npy \ --rmsd-npy orf7a_bst2_short_rmsd.npy \ --mode_chain dimer \ --indices1 0 99 \ --indices2 100 220 \ --K 6 --z-dim 8 \ --batch-size 500 \ --model-dir ./outputs/K6_z8 \ --num-cpus 64 > ./outputs/K6_z8/eval.log
Output Directory Structure
Training results are saved in the directory specified by –output-dir. The output directory contains model checkpoints, training logs, evaluation results, and summary statistics.
Output Files
Training Outputs
The following files/folders are generated during training:
- checkpoint/ Saved GMVAE model checkpoints.
- summary/: Summary statistics and evaluation results (generated after evaluation).
- train_dic.pickle: Training loss history recorded for each epoch.
- valid_dic.pickle: Validation loss history recorded for each epoch.
- train.png: Training loss versus epoch.
- validation.png: Validation loss versus epoch.
Latent Space Visualization
To visualize the learned latent representation, the pipeline generates two-dimensional t-SNE embeddings.
- learned_labels.png: Two-dimensional t-SNE projection of the latent space, colored by GMVAE cluster assignments.
- learned_rmsd.png: Two-dimensional t-SNE projection of the latent space, colored by RMSD values.
- var2d_zsel.npy: Two-dimensional t-SNE coordinates used to generate the visualization.
Clustering Results
- y_tot.npy: Predicted cluster assignment for each input conformation.
- Z_tot.npy: Learned latent-space representation for each conformation.
Evaluation Summary After running the evaluation step (–mode eval), additional summary files are generated in ./<model-dir>/summary/ including:
- latent_mean_corr_matrix.csv Correlation matrix of the latent variables.
- latent_mean_corr_matrix.png Heatmap visualization of the latent-variable correlation matrix.
- latent_mean_pca_explained.png Principal component explained variance of the latent space.
GMVAE Operation details
1. Hyperparameter Optimization (excluding K and z)
Optimize the network architecture and training parameters before performing the grid search.
Recommended Network Architecture
As a general guideline:
- 2 convolutional layers are recommended for small single-chain proteins (e.g., Trp-cage, BBA, and proteins with fewer than approximately 50 residues).
- 3 convolutional layers are recommended for larger proteins or protein dimers.
1.1. Published Hyperparameter Sets
Systems # of # of Latent # of Batch Kernel Learning # of Pooling
layers neurons dim. clusters size Temp. size rate filters sizes Ref
Trp-cage 2 64 5 8 5000 0.1 [3,3] 0.001 [64,64] [1,1] J. Chem. Phys. 155, 194108 (2021) doi: 10.1063/5.0069708
BBA 2 64 6 9 5000 0.1 [3,3] 0.001 [64,64] [2,2] J. Chem. Phys. 155, 194108 (2021) doi: 10.1063/5.0069708
villin 3 64 5 6 2500 0.05 [3,3,3] 0.001 [64,64,32] [2,2,1] J. Chem. Phys. 155, 194108 (2021) doi: 10.1063/5.0069708
BST-2/ORF7a 3 64 5 9 2500 0.05 [3,3,3] 0.001 [64,64,32] [2,2,1] BBA Biomembranes (2023)1865 (6): 184174. doi: 10.1016/j.bbamem.2023.184174.
BST-2/ORF7a 3 64 5 13 2500 0.05 [3,3,3] 0.001 [64,64,32] [2,2,1] BBA Biomembranes (2023)1865 (6): 184174. doi: 10.1016/j.bbamem.2023.184174.
ORF7b dimer 3 64 6 10 2500 0.05 [3,3,3] 0.001 [64,64,32] [2,2,1] J Phys Chem B. (2024) 128(1):150-162. doi: 10.1021/acs.jpcb.3c07105
ORF7b dimer 3 64 6 8 2500 0.05 [3,3,3] 0.001 [64,64,32] [2,2,1] J Phys Chem B. (2024) 128(1):150-162. doi: 10.1021/acs.jpcb.3c07105
BST-2/ORF7a 3 64 4 11 1000 0.05 [3,3,3] 0.001 [64,64,32] [2,2,1] J. Phys. Chem. B (2026) 130 (5): 1553–1567. doi: 10.1021/acs.jpcb.5c07529
1.2. Batch Size Recommendations
GPU memory limits the maximum batch size during training. For single-channel, three-layer GMVAE models trained on an NVIDIA V100 GPU (32 GB), the recommended batch sizes are summarized below.:
chain A chain A matrix batch size 25 25 [25x25] 5000 50 50 [50x50] 2500 100 100 [100x100] 1000 150 150 [150x150] 200
1.3. Dataset Size Considerations
The current implementation processes only complete batches during training and evaluation. If the total number of samples is not divisible by the batch size, the remaining samples in the final incomplete batch are ignored.
To maximize data utilization, select a batch size that minimizes the number of discarded samples.
# of start end data batch # of data input data frame frame included size batch leftovers 15000 0 14999 15000 1000 15 0 5739 0 5699 5700 100 57 39 5739 0 5631 5632 128 44 107 5739 0 5695 5696 64 89 43
Note: If the whole usage of data set is not a case, please specify the –start and –end in the eval mode.
1.4. Recommended Optimization Strategy
Optimize the hyperparameters in the following order:
- Batch size
- Number of convolutional layers
- Remaining network hyperparameters
- Number of clusters (K)
Based on our experience, z ≥ 4 generally provides sufficient latent-space capacity, with z = 6 serving as a good starting point for many protein systems. An initial K value between 6 and 13 is recommended, consistent with previously published GMVAE studies.
After selecting an initial configuration, perform several single-model training runs to refine the hyperparameters before proceeding to the grid search.
2. Grid Search After optimizing the training hyperparameters (excluding K and z), perform a systematic grid search over candidate values of K and z. Compare the resulting models using reconstruction loss, clustering confidence, latent-space independence, and structural interpretability to identify the optimal (K, z) configuration.
3. Evaluate the Latent Space and Select an Optimal (K, z) After completing the grid search, evaluate each model using multiple complementary metrics. The goal is to identify the lowest-complexity model that accurately represents the conformational landscape while producing confident, interpretable, and non-redundant clusters. Optimal Model Selection Criteria An optimal GMVAE model should satisfy the following criteria:
- Low reconstruction loss: Accurately reconstructs the input distance matrices while preserving conformational variability.
- Low cluster assignment entropy: Produces confident, well-separated Gaussian mixture assignments.
- Low latent multicollinearity (VIF): Latent variables are approximately independent and encode complementary structural information.
- Meaningful PCA contributions: Each latent dimension contributes measurable variance, with minimal redundancy.
- Balanced cluster populations: Avoids excessive empty or very small clusters while capturing dominant conformational states.
- Structural interpretability: Clusters correspond to physically meaningful and biologically relevant conformations.
Practical Recommendations
- Choose the smallest latent dimension (z) after reconstruction loss and cluster entropy have converged.
- Select the number of clusters (K) that balances structural diversity with sufficient samples per cluster.
- Verify that latent variables have low VIF and weak inter-latent correlations.
- Use PCA to confirm that each latent dimension contributes unique information.
- Inspect representative structures from each cluster to ensure biological interpretability.
- Prefer the smallest (K, z) configuration that meets all evaluation criteria, balancing reconstruction quality, clustering confidence, latent-space independence, cluster population, and structural interpretability. [the smallest (K, z) configuration]
data_collection.py
Required python (main) packages channels:
- conda-forge
- defaults
dependencies:
- python=3.9
- numpy
- mdtraj
- pyemma
or using provided pyemma.yml to build conda environment.
Truncating Trajectories
For large systems, it is recommended to trim the trajectory to include only the heavy atoms of interest (e.g., CA).
Data file archive: save pdb in ca_dcd/ (e.g. ca_dcd/ca_dyn0.pdb)
Example:
python data_collection.py --pdb /home/mhsieh/scratch/bst2/orf7a-cov-2_bst2_short/aa/c2_r1_cg2aa/openmm/ca_dcd/ca.pdb --traj_dirs "/home/mhsieh/scratch/bst2/orf7a-cov-2_bst2_short/aa/c2_r1_cg2aa/openmm/ca_dcd/" "/home/mhsieh/scratch/bst2/orf7a-cov-2_bst2_short/aa/c2_r2_cg2aa/openmm/ca_dcd/" --traj_prefix ca_dyn --start 501 --end 550 –-dimer –-skip 4 –-output_prefix orf7a_bst2_short
Starting Data Collection Pipeline
Length of indices1: 100 Length of indices2: 121 Feature dimension: 12100 100%|████████████████████| 100/100 [00:14<00:00, 6.82it/s] 100%|████████████████████| 100/100 [00:01<00:00, 97.15it/s] Final all_data length: 10000 Computed RMSD for 10000 frames. Saved outputs: orf7a_bst2_short_data.npy and orf7a_bst2_short_rmsd.npy
Pipeline Completed
Using tf1.15.4.yml to build conda environment (pyemma excluded). GMVAE clustering via Zaratan v100 GPU
Example of run.sh
#!/bin/bash #SBATCH --job-name=k6_z8_r1 #SBATCH -n 8 #SBATCH --partition=GPU #SBATCH --gpus=v100:1 #SBATCH --time=01:30:00 source ~/.bashrc conda init --all conda activate /home/mhsieh/scratch/tf1.15.4/ pip list | grep tensorflow echo "python gmvae_pipeline.py --mode train \ --pdb /home/mhsieh/scratch/bst2/orf7a-cov-2_bst2_short/aa/c2_r1_cg2aa/openmm/ca_dcd/ca.pdb \ --dataset-npy orf7a_bst2_short_data.npy \ --rmsd-npy orf7a_bst2_short_rmsd.npy \ --mode_chain dimer \ --indices1 0 99 \ --indices2 100 220 \ --K 6 --z-dim 8 \ --epochs 100 --batch-size 500 \ --output-dir outputs" python gmvae_pipeline.py --mode train \ --pdb /home/mhsieh/scratch/bst2/orf7a-cov-2_bst2_short/aa/c2_r1_cg2aa/openmm/ca_dcd/ca.pdb \ --dataset-npy orf7a_bst2_short_data.npy \ --rmsd-npy orf7a_bst2_short_rmsd.npy \ --mode_chain dimer \ --indices1 0 99 \ --indices2 100 220 \ --K 6 --z-dim 8 \ --epochs 100 --batch-size 500 \ --output-dir outputs
If process evaluation right after clustering, the add –evaluate_after_train option.
python gmvae_pipeline.py --mode train \ --pdb /home/mhsieh/scratch/bst2/orf7a-cov-2_bst2_short/aa/c2_r1_cg2aa/openmm/ca_dcd/ca.pdb \ --dataset-npy orf7a_bst2_short_data.npy \ --rmsd-npy orf7a_bst2_short_rmsd.npy \ --mode_chain dimer \ --indices1 0 99 \ --indices2 100 220 \ --K 6 --z-dim 4 \ --epochs 10 --batch-size 500 \ --output-dir outputs \ --evaluate_after_train
If run evaluation separately, using the following script.
#!/bin/bash #SBATCH --job-name=k11_z4_r1 #SBATCH -N 1 #SBATCH -n 64 #SBATCH --partition=standard #SBATCH --time=01:30:00 source ~/.bashrc conda init --all conda activate /home/mhsieh/scratch/tf1.15.4/ pip list | grep tensorflow echo "python gmvae_pipeline.py --mode eval \ --pdb /home/mhsieh/scratch/bst2/orf7a-cov-2_bst2_short/aa/c2_r1_cg2aa/openmm/ca_dcd/ca.pdb \ --dataset-npy orf7a_bst2_short_data.npy \ --rmsd-npy orf7a_bst2_short_rmsd.npy \ --mode_chain dimer \ --indices1 0 99 \ --indices2 100 220 \ --K 6 --z-dim 8 \ --batch-size 1 \ --output-dir outputs \ --model-dir ./outputs/K6_z8 \ --num-cpus 64" python gmvae_pipeline.py --mode eval \ --pdb /home/mhsieh/scratch/bst2/orf7a-cov-2_bst2_short/aa/c2_r1_cg2aa/openmm/ca_dcd/ca.pdb \ --dataset-npy orf7a_bst2_short_data.npy \ --rmsd-npy orf7a_bst2_short_rmsd.npy \ --mode_chain dimer \ --indices1 0 99 \ --indices2 100 220 \ --K 6 --z-dim 8 \ --batch-size 1 \ --output-dir outputs \ --model-dir ./outputs/K6_z8 \ --num-cpus 64