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
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