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.
File Dependencies 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