Minor updates to scripts.

Author: ameya98

README.md

@@ -33,6 +33,7 @@ This is because certain dependencies are tricky to install directly.
 conda create --name jamun python=3.11 -y
 conda activate jamun
 conda install -c conda-forge ambertools=23 openmm pdbfixer pyemma -y
+conda install pulchra -c bioconda -y
 ```
 
 The remaining dependencies can be installed via `pip` or [`uv`](https://docs.astral.sh/uv/getting-started/installation/) (recommended).
@@ -191,7 +192,10 @@ Please run this script with the `-h` flag to see all simulation parameters.
 ## Preprocessing
 
 ```bash
-python scripts/process_mdgen.py  --input-dir /data/bucket/kleinhej/mdgen --output-dir /data/bucket/kleinhej/mdgen/data/4AA_sims_partitioned_chunked
+source .env
+python scripts/process_mdgen.py \
+  --input-dir ${JAMUN_DATA_PATH}/mdgen \
+  --output-dir ${JAMUN_DATA_PATH}/mdgen/data/4AA_sims_partitioned_chunked
 ```
 
 ## Citation

analysis/analysis_sweep.py

@@ -11,6 +11,7 @@
 
 sys.path.append("./")
 
+from jamun.utils.slurm import wait_for_jobs
 import load_trajectory
 
 
@@ -43,65 +44,6 @@ def run_analysis(args) -> Tuple[str, Optional[str], Optional[str]]:
         return (peptide, None, None)
 
 
-def wait_for_jobs(job_ids: List[str], poll_interval: int = 60) -> int:
-    """Wait for all jobs to finish and print progress."""
-
-    previous_states = collections.defaultdict(str)
-    completion_count = 0
-    total_jobs = len(job_ids)
-    
-    while True:
-        cmd = [
-            "sacct", 
-            "-j", ",".join(job_ids),
-            "--format=JobID,State", 
-            "--noheader",
-            "--parsable2"
-        ]
-        
-        result = subprocess.run(cmd, capture_output=True, text=True)
-        current_states: Dict[str, str] = {}
-        
-        # Parse current states.
-        for line in result.stdout.strip().split('\n'):
-            if not line: continue
-            jobid, state = line.split('|')
-            if '.' not in jobid:  # Only main jobs
-                current_states[jobid] = state
-
-                # If job just completed (wasn't completed before).
-                if state == 'COMPLETED' and previous_states[jobid] != 'COMPLETED':
-                    completion_count += 1
-                    print(f"Job {jobid} completed successfully. Progress: {completion_count}/{total_jobs}")
-
-        # Update states for next iteration.
-        previous_states.update(current_states)
-        
-        # Group jobs by state for summary.
-        states_summary = collections.defaultdict(int)
-        for state in current_states.values():
-            states_summary[state] += 1
-            
-        print(f"\nStatus summary:")
-        print(f"Completed: {completion_count}/{total_jobs} ({completion_count/total_jobs*100:.1f}%)")
-        print(f"Current states: {dict(states_summary)}")
-        
-        # Check if all jobs reached terminal state.
-        all_done = all(state in ['COMPLETED', 'FAILED', 'TIMEOUT', 'OUT_OF_MEMORY', 'CANCELLED'] 
-                      for state in current_states.values())
-        
-        if all_done:
-            print("\nAll jobs finished!")
-            failures = [jid for jid, state in current_states.items() if state != 'COMPLETED']
-            if failures:
-                print(f"Failed jobs: {failures}")
-            break
-            
-        time.sleep(poll_interval)
-    
-    return completion_count
-
-
 def main():
     parser = argparse.ArgumentParser(description="Run analysis for multiple peptides")
     parser.add_argument("--csv", type=str, required=True, help="CSV file containing wandb runs")

configs/experiment/sample_mdgen.yaml

@@ -12,8 +12,7 @@ repeat_init_samples: 1
 continue_chain: true
 
 # MDGen
-wandb_train_run_path: prescient-design/jamun/lmnf3vyu
-
+wandb_train_run_path: prescient-design/jamun/brd51ln4
 
 checkpoint_type: best_so_far
 sigma: 0.04

configs/experiment/train_idrome.yaml

@@ -0,0 +1,56 @@
+# @package _global_
+
+compute_average_squared_distance_from_data: false
+
+model:
+  average_squared_distance: 0.332
+  sigma_distribution:
+    _target_: jamun.distributions.ConstantSigma
+    sigma: 0.04
+  max_radius: 1.0
+  optim:
+    lr: 0.002
+  use_torch_compile: true
+  torch_compile_kwargs:
+    fullgraph: true
+    dynamic: true
+    mode: default
+
+callbacks:
+  viz:
+    sigma_list: ["${model.sigma_distribution.sigma}"]
+
+data:
+  datamodule:
+    num_workers: 4
+    batch_size: 32
+    datasets:
+      train:
+        _target_: jamun.data.parse_datasets_from_directory
+        root: "${paths.data_path}/IDRome_v4_preprocessed/all_atom_relaxed_combined/"
+        traj_pattern: "^(.*)/traj.xtc"
+        pdb_pattern: "^(.*)/top.pdb"
+
+      val:
+        _target_: jamun.data.parse_datasets_from_directory
+        root: "${paths.data_path}/IDRome_v4_preprocessed/all_atom_relaxed_combined/"
+        traj_pattern: "^(.*)/traj.xtc"
+        pdb_pattern: "^(.*)/top.pdb"
+        subsample: 100
+
+      test:
+        _target_: jamun.data.parse_datasets_from_directory
+        root: "${paths.data_path}/IDRome_v4_preprocessed/all_atom_relaxed_combined/"
+        traj_pattern: "^(.*)/traj.xtc"
+        pdb_pattern: "^(.*)/top.pdb"
+        subsample: 100
+
+trainer:
+  val_check_interval: 30000
+  limit_val_batches: 1000
+  max_epochs: 10
+
+logger:
+  wandb:
+    group: train_idrome
+

configs/experiment/train_idrome_cg.yaml

@@ -0,0 +1,56 @@
+# @package _global_
+
+compute_average_squared_distance_from_data: false
+
+model:
+  average_squared_distance: 0.332
+  sigma_distribution:
+    _target_: jamun.distributions.ConstantSigma
+    sigma: 0.08
+  max_radius: 2.0
+  optim:
+    lr: 0.002
+  use_torch_compile: true
+  torch_compile_kwargs:
+    fullgraph: true
+    dynamic: true
+    mode: default
+
+callbacks:
+  viz:
+    sigma_list: ["${model.sigma_distribution.sigma}"]
+
+data:
+  datamodule:
+    num_workers: 4
+    batch_size: 32
+    datasets:
+      train:
+        _target_: jamun.data.parse_datasets_from_directory
+        root: "${paths.data_path}/IDRome_v4_preprocessed/flat/"
+        traj_pattern: "^(.*)/traj.xtc"
+        pdb_pattern: "^(.*)/top.pdb"
+
+      val:
+        _target_: jamun.data.parse_datasets_from_directory
+        root: "${paths.data_path}/IDRome_v4_preprocessed/flat/"
+        traj_pattern: "^(.*)/traj.xtc"
+        pdb_pattern: "^(.*)/top.pdb"
+        subsample: 100
+
+      test:
+        _target_: jamun.data.parse_datasets_from_directory
+        root: "${paths.data_path}/IDRome_v4_preprocessed/flat/"
+        traj_pattern: "^(.*)/traj.xtc"
+        pdb_pattern: "^(.*)/top.pdb"
+        subsample: 100
+
+trainer:
+  val_check_interval: 30000
+  limit_val_batches: 1000
+  max_epochs: 10
+
+logger:
+  wandb:
+    group: train_idrome
+

configs/experiment/train_mdgen.yaml

@@ -32,6 +32,7 @@ data:
         pdb_pattern: "^(....).pdb"
         as_iterable: true
         subsample: 5
+        start_at_random_frame: true
 
       val:
         _target_: jamun.data.parse_datasets_from_directory
@@ -40,6 +41,7 @@ data:
         pdb_pattern: "^(....).pdb"
         as_iterable: true
         subsample: 100
+        start_at_random_frame: true
 
       test:
         _target_: jamun.data.parse_datasets_from_directory
@@ -48,8 +50,11 @@ data:
         pdb_pattern: "^(....).pdb"
         as_iterable: true
         subsample: 100
+        start_at_random_frame: true
 
 trainer:
+  val_check_interval: 30000
+  limit_val_batches: 1000
   max_epochs: 10
 
 logger:

scripts/IDRome/README.md

@@ -8,13 +8,26 @@ conda install pulchra -c bioconda --yes
 ```
 
 ```bash
-python scripts/generate_data/run_simulation.py /homefs/home/daigavaa/jamun/145_181/all_atom/top_AA.pdb --energy-minimization-only --energy-minimization-steps=5000
+source .env
+sbatch scripts/IDRome/to_all_atom_batched.sh \
+    ${JAMUN_DATA_PATH}/IDRome_v4_preprocessed/flat \
+    ${JAMUN_DATA_PATH}/IDRome_v4_preprocessed/flat_by_frame/ \
+    ${JAMUN_DATA_PATH}/IDRome_v4_preprocessed/all_atom/ \
+    1000
 ```
 
 ```bash
-sbatch to_all_atom_batched.sh /data/bucket/kleinhej/IDRome_v4_preprocessed/flat /data/bucket/kleinhej/IDRome_v4_preprocessed/flat_by_frame/ /data/bucket/kleinhej/IDRome_v4_preprocessed/all_atom/ 1000
+source .env
+sbatch scripts/IDRome/relax_structures_batched.sh \
+    ${JAMUN_DATA_PATH}/IDRome_v4_preprocessed/all_atom \
+    ${JAMUN_DATA_PATH}/IDRome_v4_preprocessed/all_atom_relaxed \
+    1000
 ```
 
 ```bash
-sbatch relax_structures_batched.sh /data/bucket/kleinhej/IDRome_v4_preprocessed/all_atom /data/bucket/kleinhej/IDRome_v4_preprocessed/all_atom_relaxed 1000
+source .env
+sbatch scripts/IDRome/combine_frames.sh \
+    ${JAMUN_DATA_PATH}/IDRome_v4_preprocessed/all_atom_relaxed/ \
+    ${JAMUN_DATA_PATH}/IDRome_v4_preprocessed/flat/ \
+    ${JAMUN_DATA_PATH}/IDRome_v4_preprocessed/all_atom_relaxed_combined/
 ```

scripts/IDRome/combine_frames.py

@@ -12,9 +12,8 @@
 
 
 
-def combine_frames(args, use_srun: bool = True) -> None:
+def combine_frames(name: str, input_dir: str, original_traj_dir: str, output_dir: str) -> None:
     """Combine relaxed IDRome v4 all-atom frames."""
-    name, input_dir, original_traj_dir, output_dir = args
 
     traj_AA = None
     frames = sorted(os.listdir(os.path.join(input_dir, name)),
@@ -39,33 +38,25 @@ def combine_frames(args, use_srun: bool = True) -> None:
             top_AA.add_atom(atom.name, element=atom.element, residue=res)
     top_AA.create_standard_bonds()
 
-    traj = md.load_xtc(os.path.join(original_traj_dir, f'{name}.xtc'), top=os.path.join(original_traj_dir, f'{name}.pdb'))
-    traj_AA = md.Trajectory(traj_AA.xyz, top_AA, traj.time, traj.unitcell_lengths, traj.unitcell_angles)
-    traj_AA[0].save_pdb(os.path.join(output_dir, f"{name}.pdb"))
-    traj_AA.save_xtc(os.path.join(output_dir, f"{name}.xtc"))
+    original_traj_path = os.path.join(original_traj_dir, name, 'traj.xtc')
+    original_top_path = os.path.join(original_traj_dir, name, 'top.pdb')
+    original_traj = md.load_xtc(original_traj_path, top=original_top_path)
+    original_traj = original_traj[0:traj_AA.n_frames]
 
+    os.makedirs(os.path.join(output_dir, name), exist_ok=True)
+    traj_AA = md.Trajectory(traj_AA.xyz, top_AA, original_traj.time, original_traj.unitcell_lengths, original_traj.unitcell_angles)
+    traj_AA[0].save_pdb(os.path.join(output_dir, name, 'top.pdb'))
+    traj_AA.save_xtc(os.path.join(output_dir, name, 'traj.xtc'))
+
+    py_logger.info(f"Successfully processed {name}")
 
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(description='Convert IDRome v4 data to all-atom.')
+    parser.add_argument('--name', help='Name of the trajectory.', type=str, required=True)
     parser.add_argument('--input-dir', help='Directory of relaxed all-atom trajectories (stored in each folder).', type=str, required=True)
     parser.add_argument('--original-traj-dir', help='Directory of original coarse-grained trajectories (stored in each folder).', type=str, required=True)
     parser.add_argument('--output-dir', '-o', help='Output directory to save combined relaxed all-atom trajectories (stored in each folder).', type=str, required=True)
-    parser.add_argument('--num-workers', type=int, default=multiprocessing.cpu_count(),
-                      help='Number of parallel workers')
     args = parser.parse_args()
 
-    # Run in parallel.
-    names = sorted(os.listdir(args.input_dir))
-    preprocess_args = list(
-        zip(
-            names,
-            [args.input_dir] * len(names),
-            [args.original_traj_dir] * len(names),
-            [args.output_dir] * len(names),
-        )
-    )
-    with ProcessPoolExecutor(max_workers=args.num_workers) as executor:
-        results = list(executor.map(combine_frames, preprocess_args))
-
-
+    combine_frames(args.name, args.input_dir, args.original_traj_dir, args.output_dir)

scripts/IDRome/combine_frames.sh

@@ -0,0 +1,51 @@
+#!/bin/bash
+#SBATCH --partition=cpu
+#SBATCH --mem=1G
+#SBATCH --cpus-per-task=2
+#SBATCH --job-name=combine_frames
+#SBATCH --output=logs/%j_combine_frames.log
+#SBATCH --error=logs/%j_combine_frames.err
+#SBATCH --array=0-1
+
+# Directory containing all input directories
+BASE_INPUT_DIR="$1"
+# Directory containing all original coarse-grained directories
+BASE_ORIGINAL_DIR="$2"
+# Directory to store output
+BASE_OUTPUT_DIR="$3"
+
+eval "$(conda shell.bash hook)"
+conda activate jamun
+
+# Get list of all directories and store in an array
+# You can use a file with directory names or generate the list dynamically
+DIRECTORIES=($(ls -d ${BASE_INPUT_DIR}/*/ | sort | xargs -n 1 basename))
+
+# Each job processes 50 directories
+START_IDX=$((SLURM_ARRAY_TASK_ID * 50))
+END_IDX=$(( (SLURM_ARRAY_TASK_ID + 1) * 50 - 1 ))
+
+for DIR_INDEX in $(seq ${START_IDX} ${END_IDX}); do
+    NAME="${DIRECTORIES[${DIR_INDEX}]}"
+
+    echo "Processing directory: ${NAME} (index: ${DIR_INDEX})"
+
+    # Create output directory
+    mkdir -p "${BASE_OUTPUT_DIR}/${NAME}"
+
+    # Check if the input frame exists
+    if [ ! -f "${BASE_INPUT_DIR}/${NAME}/0_minimized_protein_0.pdb" ]; then
+        echo "Input frame 0 does not exist in ${NAME}. Skipping."
+        continue
+    fi
+
+    python scripts/IDRome/combine_frames.py \
+        --name "${NAME}" \
+        --input-dir "${BASE_INPUT_DIR}" \
+        --original-traj-dir "${BASE_ORIGINAL_DIR}" \
+        --output-dir "${BASE_OUTPUT_DIR}"
+            
+    echo "Completed processing directory ${NAME}"
+done
+
+exit 0