Skip to content

README.md

Latest commit

 

History

History
327 lines (253 loc) · 9.55 KB

README.md

File metadata and controls

327 lines (253 loc) · 9.55 KB
Raw

CGML Base Model Repository

This repository contains the CGSchnet ML model for molecular dynamics and machine learning pipelines, providing tools for training, testing, and building upon coarse-grained neural network force fields. The model architecture is inspired by CGSchNet.

Overview

This repository implements a machine learning framework for coarse-grained molecular dynamics simulations. The pipeline consists of three main stages:

  1. Preprocessing: Convert all-atom MD trajectories into coarse-grained representations and fit classical prior force fields
  2. Training: Train neural network models to predict delta forces (corrections to the prior force field)
  3. Simulation: Run MD simulations using the trained models combined with prior force fields

The system supports multiple coarse-graining schemes (CA, CACB) and various prior force field configurations with bonds, angles, dihedrals, and non-bonded terms.

Installation

⚠️ Note: Installation instructions will be provided as the code is finalized for public release.

Prerequisites

  • Python 3.8+
  • PyTorch 2.0+
  • CUDA-capable GPU (recommended)
  • Required Python packages: torchmd, moleculekit, mdtraj, numpy, pyyaml

Model Inputs

Training Configuration

Key hyperparameters specified in config files (e.g., configs/config.yaml):

  • embedding_dimension: Dimension of learned atomic embeddings
  • num_layers: Number of interaction layers in the neural network
  • num_rbf: Number of radial basis functions
  • cutoff: Interaction cutoff distance
  • max_num_neighbors: Maximum neighbors per atom
  • activation: Activation function (e.g., 'silu')

Training Arguments

python train.py <input_dir> <output_dir> [options]

Required:

  • input_dir: Path to preprocessed data directory
  • output_dir: Directory for saving checkpoints and results

Key Options:

  • --config: Path to model configuration YAML file
  • --gpus: GPU IDs to use (e.g., "0,1,2,3" or "cpu")
  • --batch: Batch size for training
  • --epochs: Number of training epochs
  • --lr: Learning rate
  • --wd: Weight decay for regularization
  • --atoms-per-call: Max atoms per sub-batch (for memory management)
  • --energy-weight: Weight for energy loss term
  • --force-weight: Weight for force loss term
  • --val-ratio: Validation set ratio (default: 0.1)
  • --early-stopping: Patience for early stopping (default: 1)

Learning Rate Schedulers:

  • --exp-lr: Exponential decay (parameter: gamma)
  • --cos-lr: Cosine annealing (parameters: T_max, eta_min)
  • --plateau-lr: Reduce on plateau (parameters: factor, patience, threshold, min_lr)

Preprocessing

Preprocessing converts all-atom MD trajectories into coarse-grained representations and fits prior force fields.

Basic Usage

python preprocess.py <input_path> -o <output_dir> --prior <prior_type>

Example: Preprocessing with CA Representation

python preprocess.py benchmark_set.yaml \
    -o /path/to/output/preprocessed_CA_Majewski \
    --prior CA_Majewski2022_v1 \
    --temp 300 \
    --num-cores 32

Preprocessing Arguments

Required:

  • input: Input directory or YAML config file listing trajectories
  • -o, --output: Output directory for preprocessed data
  • --prior: Prior force field type (see Available Priors below)

Optional:

  • --pdbids: Specific PDB IDs to process
  • --num-frames: Number of frames to process per trajectory
  • --frame-slice: Python slice notation (e.g., "0:1000:10")
  • --temp: Temperature in Kelvin (default: 300)
  • --optimize-forces: Use statistically optimal force aggregation
  • --prior-file: Use existing prior instead of fitting new one
  • --no-box: Disable periodic boundary conditions
  • --num-cores: Number of parallel processes (default: 32)

Available Prior Force Fields

  • CA: Basic CA representation with bonds only
  • CA_lj: CA with bonds and repulsion
  • CA_lj_angle: CA with bonds, angles, and repulsion
  • CA_lj_angle_dihedral: CA with all classical terms
  • CA_lj_angleXCX_dihedralX: CA with unified angle/dihedral parameters
  • CA_Majewski2022_v1: Parameters from Majewski et al. 2022
  • CACB: CA-CB representation
  • CACB_lj: CA-CB with repulsion
  • CACB_lj_angle_dihedral: CA-CB with all terms

Input Data Format

Preprocessing expects HDF5 files with structure:

input_dir/
├── pdb_id_1/
│   └── result/
│       └── output_pdb_id_1.h5  # Contains coordinates, forces, box (optional)
├── pdb_id_2/
│   └── result/
│       └── output_pdb_id_2.h5
...

Output Structure

output_dir/
├── priors.yaml              # Fitted prior force field parameters
├── prior_params.json        # Force field configuration
├── result/
│   ├── info.json           # Preprocessing metadata
│   └── ok_list.txt         # Successfully processed PDB IDs
└── <pdb_id>/
    ├── raw/
    │   ├── coordinates.npy  # CG coordinates
    │   ├── forces.npy       # CG forces from all-atom MD
    │   ├── deltaforces.npy  # Delta forces (MD - prior)
    │   ├── embeddings.npy   # Atom type embeddings
    │   └── box.npy         # Periodic box (if used)
    └── processed/
        └── topology.psf     # CG topology

Training

Train neural network models to predict corrections to the prior force field.

Basic Training Example

python train.py \
    /path/to/preprocessed_data \
    /path/to/results \
    --config configs/config.yaml \
    --gpus 0,1,2,3 \
    --batch 4 \
    --epochs 35 \
    --lr 0.001 \
    --exp-lr 0.85 \
    --atoms-per-call 140000 \
    --energy-weight 0.1 \
    --force-weight 1.0

Training with Energy Matching

python train.py \
    /path/to/preprocessed_with_tica \
    /path/to/results_energy_matching \
    --config configs/config_cutoff2.yaml \
    --gpus 0,1,2,3 \
    --batch 4 \
    --epochs 35 \
    --lr 0.001 \
    --wd 0 \
    --exp-lr 0.85 \
    --atoms-per-call 140000 \
    --energy-weight 0.5 \
    --force-weight 1.0

Resuming Training

To resume from a checkpoint, simply provide the same output directory:

python train.py \
    /path/to/preprocessed_data \
    /path/to/existing_results \
    --config configs/config.yaml \
    --gpus 0,1,2,3

Training Outputs

results_dir/
├── checkpoint.pth           # Latest model checkpoint
├── checkpoint-best.pth      # Best validation loss checkpoint
├── history.npy             # Training/validation loss history
├── epoch_history.json      # Detailed per-epoch metrics
├── training_info.json      # Training configuration
├── priors.yaml            # Copy of prior force field
└── prior_params.json      # Copy of prior configuration

Running Simulations

Use trained models to run coarse-grained MD simulations.

Basic Simulation

python simulate.py \
    /path/to/results/checkpoint.pth \
    /path/to/starting_structure.pdb \
    --output simulation.h5 \
    --steps 1000000 \
    --save-steps 1000 \
    --temperature 300 \
    --timestep 2 \
    --replicas 1

Multiple Replicas

python simulate.py \
    /path/to/results/checkpoint-best.pth \
    /path/to/structure1.pdb /path/to/structure2.pdb /path/to/structure3.pdb \
    --output replicas.pkl \
    --steps 5000000 \
    --save-steps 1000 \
    --temperature 300 \
    --replicas 3

Prior-Only Simulation

To run simulations using only the classical prior (without ML corrections):

python simulate.py \
    /path/to/results/checkpoint.pth \
    /path/to/starting_structure.pdb \
    --output prior_only.h5 \
    --prior-only \
    --steps 1000000 \
    --save-steps 1000

Simulation Arguments

Required:

  • checkpoint_path: Path to trained model checkpoint
  • processed_path: One or more starting structures (PDB or preprocessed directories)

Key Options:

  • --output: Output trajectory file (.pdb, .h5, or .pkl)
  • --steps: Total simulation steps
  • --save-steps: Save frequency (frames)
  • --temperature: Simulation temperature in Kelvin (default: 300)
  • --timestep: Integration timestep in femtoseconds (default: 1)
  • --replicas: Number of parallel replicas (default: 1)
  • --prior-only: Run without ML model (classical prior only)
  • --no-box: Disable periodic boundary conditions
  • --max-num-neighbors: Override max neighbors parameter

Output Formats

  • PDB format (.pdb): Human-readable, compatible with visualization tools
  • HDF5 format (.h5): Compact binary format with energies and metadata
  • Pickle format (.pkl): Contains MDTraj trajectory objects for analysis

Utilities

Editing Checkpoints

Reset optimizer or scheduler state in checkpoints:

python edit_checkpoint.py /path/to/checkpoint.pth --reset-optimizer --reset-scheduler

View checkpoint information:

python edit_checkpoint.py /path/to/checkpoint.pth --info

Learning Prior Force Fields

Train classical prior parameters directly from data:

python learn_prior.py \
    /path/to/preprocessed_data \
    /path/to/prior_output \
    --epochs 50 \
    --batch 10 \
    --lr 0.005

Inspired by CGSchNet.

License

License information will be added soon.

Contributing

We welcome contributions! Please use GitHub Issues to:

  • Request new features
  • Report bugs
  • Suggest improvements
  • Ask questions about the pipeline

Status

⚠️ Note: Code is currently being prepared for public release. Full documentation and tutorials will be available within 1-2 weeks.