## data collection ##

python data_collection_v4.py --help
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]

Dimerization data preprocessing pipeline.

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. 1 441.
  --indices2 INDICES2 INDICES2
                        Range for chain 2 MDTraj atom indices, e.g. 442 882.
  --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.


Example 1:
Explicit trajectory list:
   --traj_files traj1.dcd traj2.dcd traj3.dcd

Example 2: 
Directory + prefix + index range:
   --traj_dirs replica1 replica2 --traj_ext dcd --traj_prefix ca_dyn --start 501 --end 505

(in Example 2, the expected trajectory list:
replica1/ca_dyn501.dcd 
replica1/ca_dyn502.dcd 
replica1/ca_dyn503.dcd 
replica1/ca_dyn504.dcd 
replica1/ca_dyn505.dcd 
replica2/ca_dyn501.dcd 
replica2/ca_dyn502.dcd 
replica2/ca_dyn503.dcd 
replica2/ca_dyn504.dcd 
replica2/ca_dyn505.dcd 
)

The example 2 is suggested if the trajectory file are systemical named. Please see scripts 
scripts/run_datacoll_list.csh and scripts/run_datacoll.csh for details.

## gmvae ##

usage: gmvae_pipeline_v2.py [-h] --mode {train,grid,eval} --pdb PDB
                            --dataset-npy DATASET_NPY --rmsd-npy RMSD_NPY
                            [--mode_chain {single_chain,dimer}]
                            [--indices1 INDICES1 INDICES1]
                            [--indices2 INDICES2 INDICES2]
                            [--batch-size BATCH_SIZE] [--lr LR]
                            [--epochs EPOCHS] [--num-layers NUM_LAYERS]
                            [--z-dim Z_DIM] [--temp TEMP] [--K K]
                            [--restore RESTORE] [--dense-n DENSE_N]
                            [--filters FILTERS [FILTERS ...]]
                            [--kernel-sizes KERNEL_SIZES [KERNEL_SIZES ...]]
                            [--pool-sizes POOL_SIZES [POOL_SIZES ...]]
                            [--valid-ratio VALID_RATIO] [--dropout DROPOUT]
                            [--r-nent R_NENT] [--r-label R_LABEL]
                            [--output-dir OUTPUT_DIR]
                            [--grid-ks [GRID_KS [GRID_KS ...]]]
                            [--grid-zs [GRID_ZS [GRID_ZS ...]]]
                            [--evaluate-after-train] [--num-cpus NUM_CPUS]
                            [--model-dir MODEL_DIR]

GMVAE training & grid testing pipeline

optional arguments:
  -h, --help            show this help message and exit
  --mode {train,grid,eval}
                        Mode to run: train | grid | eval
  --pdb PDB             Path to pdb file
  --dataset-npy DATASET_NPY
                        Path to dataset .npy file
  --rmsd-npy RMSD_NPY   Path to RMSD .npy file
  --mode_chain {single_chain,dimer}
                        Single chain or dimer analysis (affects indices_range2
                        selection)
  --indices1 INDICES1 INDICES1
                        Range for chain 1 indices, e.g., 0 99.
  --indices2 INDICES2 INDICES2
                        Range for chain 2 indices, e.g., 100 220.
  --batch-size BATCH_SIZE
  --lr LR
  --epochs EPOCHS
  --num-layers NUM_LAYERS
  --z-dim Z_DIM
  --temp TEMP
  --K K
  --restore RESTORE
  --dense-n DENSE_N
  --filters FILTERS [FILTERS ...]
  --kernel-sizes KERNEL_SIZES [KERNEL_SIZES ...]
  --pool-sizes POOL_SIZES [POOL_SIZES ...]
  --valid-ratio VALID_RATIO
  --dropout DROPOUT
  --r-nent R_NENT
  --r-label R_LABEL
  --output-dir OUTPUT_DIR
                        Base output directory
  --grid-ks [GRID_KS [GRID_KS ...]]
                        List of K values for grid
  --grid-zs [GRID_ZS [GRID_ZS ...]]
                        List of z-dim values for grid
  --evaluate-after-train
                        Run evaluation after each training run (CPU)
  --num-cpus NUM_CPUS   Number of CPU cores to pass to evaluation
  --model-dir MODEL_DIR
                        Model output dir to evaluate (used in eval mode)


## contact map & top structures ##
To generate top structures, please use --mode trajectory and specify --generate_top_structures.
To ensure valid results, please input the exact same trajectories as the data collection.

python src/utils/contact_heatmap_top_structures.py -h
usage: contact_heatmap_top_structures.py [-h] --mode {dataset,trajectory}
                                         --mode_chain {single_chain,dimer}
                                         [--A_dataset A_DATASET]
                                         [--B_dataset B_DATASET]
                                         [--AB_dataset AB_DATASET]
                                         [--A_size A_SIZE] [--B_size B_SIZE]
                                         [--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]
                                         [--skip SKIP]
                                         [--skip_first_frame SKIP_FIRST_FRAME]
                                         [--skip_after_frame SKIP_AFTER_FRAME]
                                         [--selected_atom SELECTED_ATOM]
                                         [--A_indices A_INDICES A_INDICES]
                                         [--B_indices B_INDICES B_INDICES]
                                         --cluster_folder CLUSTER_FOLDER
                                         [--threshold THRESHOLD]
                                         [--cutoff CUTOFF]
                                         [--output_dir OUTPUT_DIR]
                                         [--save_generated_dataset]
                                         [--generate_top_structures]
                                         [--top_n TOP_N]
                                         [--mean_iterations MEAN_ITERATIONS]
                                         [--representative_output_dir REPRESENTATIVE_OUTPUT_DIR]

Generate cluster-specific contact heatmaps.

optional arguments:
  -h, --help            show this help message and exit
  --mode {dataset,trajectory}
                        Use existing .npy distance datasets or generate them
                        from trajectories.
  --mode_chain {single_chain,dimer}
                        single_chain: only Chain A self-contact map; dimer:
                        Chain A self, Chain B self, and A-B inter-chain maps.
  --A_dataset A_DATASET
                        Path to Chain A self-contact distance dataset .npy
  --B_dataset B_DATASET
                        Path to Chain B self-contact distance dataset .npy
  --AB_dataset AB_DATASET
                        Path to inter-chain A-B distance dataset .npy
  --A_size A_SIZE       Number of selected atoms/beads in Chain A dataset.
  --B_size B_SIZE       Number of selected atoms/beads in Chain B dataset.
  --pdb PDB             Topology file (.pdb or .gro). Required for trajectory
                        mode.
  --traj_files TRAJ_FILES [TRAJ_FILES ...]
                        Explicit list of trajectory files.
  --traj_dirs TRAJ_DIRS [TRAJ_DIRS ...]
                        Directories containing trajectories.
  --traj_prefix TRAJ_PREFIX
                        Trajectory prefix when using --traj_dirs.
  --traj_ext {dcd,trr,xtc}
                        Trajectory extension.
  --start START         Starting trajectory index.
  --end END             Ending trajectory index.
  --skip SKIP           Frame stride during trajectory loading.
  --skip_first_frame SKIP_FIRST_FRAME
                        Discard first N frames after loading.
  --skip_after_frame SKIP_AFTER_FRAME
                        Last frame index to keep after loading and
                        concatenation (inclusive). Default: keep all remaining
                        frames.
  --selected_atom SELECTED_ATOM
                        Atom/bead name, e.g. CA for AA or BB for CG.
  --A_indices A_INDICES A_INDICES
                        MDTraj atom-index range for chain A, e.g. 0 441
  --B_indices B_INDICES B_INDICES
                        MDTraj atom-index range for chain B, e.g. 442 882
  --cluster_folder CLUSTER_FOLDER
                        Folder containing y_tot.npy
  --threshold THRESHOLD
                        Cluster assignment confidence threshold.
  --cutoff CUTOFF       Contact cutoff in nm.
  --output_dir OUTPUT_DIR
                        Output directory.
  --save_generated_dataset
                        Save generated trajectory datasets in output_dir.
  --generate_top_structures
                        In trajectory mode, identify structures closest to the
                        aligned mean coordinates of each cluster and save
                        representative PDB/DCD files.
  --top_n TOP_N         Number of lowest-RMSD structures to save per cluster
                        (default: 50).
  --mean_iterations MEAN_ITERATIONS
                        Number of align-average iterations used to construct
                        each cluster mean.
  --representative_output_dir REPRESENTATIVE_OUTPUT_DIR
                        Output directory for representative structures.
                        Default: <output_dir>/representative_structures.

