conditional gen

.gitignore

@@ -24,3 +24,5 @@ torch_compile_debug
 *.profile*
 **/*.log
 **/*.err
+scripts/study
+wandb/*

configs/experiment/sample_capped_2AA.yaml

@@ -1,36 +1,37 @@
 # @package _global_
 
-defaults:
-  - override /callbacks:
-    - sampler/save_trajectory.yaml
-  - _self_
+# defaults:
+#   - override /callbacks:
+#     - sampler/save_trajectory.yaml
+#   - _self_
 
 
 init_datasets:
   _target_: jamun.data.parse_datasets_from_directory
-  root: "${paths.data_path}/capped_diamines/timewarp_splits/test"
+  root: "${paths.data_path}/capped_diamines/timewarp_splits/train"
   traj_pattern: "^(.*).xtc"
   pdb_pattern: "^(.*).pdb"
   subsample: 1
-  num_frames: 60000
+  filter_codes: ['ALA_ALA']
+  num_frames: 320000
 
 
-num_sampling_steps_per_batch: 20000
-num_batches: 5
-num_init_samples_per_dataset: 1
+num_sampling_steps_per_batch: 1000
+num_batches: 10
+num_init_samples_per_dataset: 50
 repeat_init_samples: 1
 continue_chain: true
 
 # New 2AA
-wandb_train_run_path: prescient-design/jamun/yfz3vpzg
+wandb_train_run_path: sule-shashank/jamun/370wpt17
 
 checkpoint_type: best_so_far
 sigma: 0.04
 M: 1.0
-delta: 0.04
+delta: ${sigma}
 friction: 1.0
 inverse_temperature: 1.0
-score_fn_clip: 100.0
+score_fn_clip: null
 
 sampler:
   _target_: jamun.sampling.Sampler
@@ -39,3 +40,4 @@ sampler:
 logger:
   wandb:
     group: sample_capped_2AA
+    tags: ['ALA_ALA', 'sigma_0.04', 'standard JAMUN']

configs/experiment/sample_capped_single_shape_conditioning.yaml

@@ -0,0 +1,51 @@
+# @package _global_
+
+defaults: 
+  - override /model: denoiser_conditional_pretrained.yaml
+# defaults:
+#  - override /callbacks:
+#     - sampler/save_trajectory.yaml
+#  - _self_
+
+# callbacks:
+#   viz:
+#     sigma_list: ["${model.sigma_distribution.sigma}"]
+
+init_datasets:
+  _target_: jamun.data.parse_datasets_from_directory
+  root: "${paths.data_path}/capped_diamines/timewarp_splits/train"
+  traj_pattern: "^(.*).xtc"
+  pdb_pattern: "^(.*).pdb"
+  filter_codes: ['ALA_ALA']
+  as_iterable: false
+  subsample: 100
+  total_lag_time: 10
+  lag_subsample_rate: 100
+
+num_sampling_steps_per_batch: 1000
+num_batches: 10
+num_init_samples_per_dataset: 10
+repeat_init_samples: 1
+continue_chain: false
+
+# Add your wandb run path here
+wandb_train_run_path: sule-shashank/jamun/jqp09yv1
+# checkpoint_dir: /data2/sules/jamun-conditional-runs/outputs/train/dev/runs/2025-07-01_16-15-23/checkpoints
+# checkpoint_dir: /data2/sules/jamun-conditional-runs/old_outputs/outputs/train/dev/runs/2a381c2d310e3d0789831338/checkpoints
+checkpoint_type: last
+
+sigma: 0.01
+M: 1.0
+delta: ${sigma}
+friction: 1.0
+inverse_temperature: 1.0
+score_fn_clip: 100.0
+
+sampler:
+  _target_: jamun.sampling.SamplerMemory
+  devices: 1
+
+logger:
+  wandb:
+    group: sample_ALA_ALA_conditional
+    notes: stellar-sweep-25

configs/experiment/sample_custom.yaml

@@ -1,24 +1,35 @@
 # @package _global_
 
-num_sampling_steps_per_batch: 1000
-num_batches: 1
-num_init_samples_per_dataset: 1
-repeat_init_samples: 1
-continue_chain: true
+name: cfg  # This ensures the config is saved with the 'cfg' key
 
-# wandb_train_run_path: prescient-design/jamun/zzt8s3rc
+# defaults:
+#  - override /callbacks:
+#     - sampler/save_trajectory.yaml
+#  - _self_
 
-# Old 4AA
-# wandb_train_run_path: prescient-design/jamun/ibtxmwcr
+# callbacks:
+#   viz:
+#     sigma_list: ["${model.sigma_distribution.sigma}"]
 
-# New 4AA
-wandb_train_run_path: prescient-design/jamun/6297yugb
+init_datasets:
+  _target_: jamun.data.parse_datasets_from_directory
+  root: "${paths.data_path}/timewarp/2AA-1-large/train/"
+  traj_pattern: "^(.*)-traj-arrays.npz"
+  pdb_file: AA-traj-state0.pdb
+  filter_codes:
+    - AA
+  subsample: 1
+  max_datasets: 1
 
-# Finetuned new 4AA
-# wandb_train_run_path: prescient-design/jamun/x6rwt91k
+num_sampling_steps_per_batch: 1000
+num_batches: 100
+num_init_samples_per_dataset: 1
+repeat_init_samples: 1
+continue_chain: true
 
-# init_pdb: /data/bucket/kleinhej/fast-folding/processed/chignolin/filtered.pdb
-init_pdbs: ???
+# Add your wandb run path here
+wandb_train_run_path: sule-shashank/jamun/y4rm5488
+# checkpoint_dir: outputs/train/dev/runs/2025-06-11_20-16-04/wandb/latest-run/checkpoints
 
 checkpoint_type: best_so_far
 sigma: 0.04
@@ -28,13 +39,9 @@ friction: 1.0
 inverse_temperature: 1.0
 score_fn_clip: 100.0
 
-init_datasets:
-  _target_: jamun.data.create_dataset_from_pdbs
-  pdbfiles: ${init_pdbs}
-
-finetune_on_init:
-  num_steps: ???
-  batch_size: 16
+# Model configuration
+model:
+  _target_: scratch.denoiser_test.Denoiser.load_from_checkpoint
 
 sampler:
   _target_: jamun.sampling.Sampler
@@ -43,3 +50,4 @@ sampler:
 logger:
   wandb:
     group: sample_custom
+    notes: Custom sampling with denoiser, conditioner 

configs/experiment/sample_enhanced_conditioning.yaml

@@ -0,0 +1,106 @@
+# @package _global_
+
+defaults: 
+  - override /model: denoiser_conditional_pretrained.yaml
+  - override /callbacks: null
+
+# init_datasets:
+#   _target_: jamun.data.parse_datasets_from_directory
+#   root: "${paths.data_path}/capped_diamines/timewarp_splits/train"
+#   traj_pattern: "^(.*).xtc"
+#   pdb_pattern: "^(.*).pdb"
+#   filter_codes: ['ALA_ALA']
+#   as_iterable: false
+#   subsample: 10
+#   total_lag_time: 5
+#   lag_subsample_rate: 1
+#   max_datasets: 1
+#   label_override: "ALA_ALA"
+
+init_datasets:
+  _target_: jamun.data.parse_datasets_from_directory
+  root: "/data2/sules/ALA_ALA_enhanced_full_swarm/val"
+  traj_pattern: "^(.*).xtc"
+  pdb_pattern: "^(.*).pdb"
+  as_iterable: false
+  subsample: 1
+  total_lag_time: 5
+  lag_subsample_rate: 1
+  max_datasets: 10
+  label_override: "ALA_ALA"
+
+# model:
+#   conditioner:
+#     N_structures: ${init_datasets.total_lag_time}
+
+num_sampling_steps_per_batch: 1000
+num_batches: 10
+num_init_samples_per_dataset: 1
+repeat_init_samples: 1
+continue_chain: true
+
+# Add your checkpoint path here - update with actual trained model path
+wandb_train_run_path: "sule-shashank/jamun/qiutegoj"
+# checkpoint_type: last
+# checkpoint_dir: "/data2/sules/jamun-conditional-runs/outputs/train/dev/runs/2025-07-31_00-43-14/checkpoints/"
+checkpoint_type: last
+
+sigma: 0.06
+M: 1.0
+delta: 0.066
+friction: 1.2
+inverse_temperature: 1.0
+score_fn_clip: null
+
+sampler:
+  _target_: jamun.sampling.SamplerMemory
+  devices: 1
+
+# Evaluation dataset - standard ALA_ALA from capped diamines for computing metrics
+eval_dataset:
+  _target_: jamun.data.parse_datasets_from_directory
+  root: "${paths.data_path}/capped_diamines/timewarp_splits/train"
+  traj_pattern: "^(.*).xtc"
+  pdb_pattern: "^(.*).pdb"
+  filter_codes: ['ALA_ALA']
+  as_iterable: false
+  subsample: 10
+  total_lag_time: 5
+  lag_subsample_rate: 1
+  max_datasets: 10
+  label_override: "ALA_ALA"
+
+# Override ALL callbacks to use eval_dataset for metrics computation
+callbacks:
+  measure_sampling_time:
+    _target_: jamun.callbacks.sampler.MeasureSamplingTimeCallback
+  chemical_validity:
+    _target_: jamun.callbacks.sampler.ChemicalValidityMetricsCallback
+    datasets: ${eval_dataset}
+    bond_length_tolerance: 0.2
+    volume_exclusion_tolerance: 0.1
+    num_molecules_per_trajectory: 100
+  ramachandran_plot:
+    _target_: jamun.callbacks.sampler.RamachandranPlotMetricsCallback
+    datasets: ${eval_dataset}
+  trajectory_visualizer:
+    _target_: jamun.callbacks.sampler.TrajectoryVisualizerCallback
+    datasets: ${eval_dataset}
+    num_frames_to_animate: 100
+  sample_visualizer:
+    _target_: jamun.callbacks.sampler.SampleVisualizerCallback
+    datasets: ${eval_dataset}
+    num_samples_to_plot: 16
+    subsample: 100
+  score_distribution:
+    _target_: jamun.callbacks.sampler.ScoreDistributionCallback
+    datasets: ${eval_dataset}
+  save_trajectory:
+    _target_: jamun.callbacks.sampler.SaveTrajectoryCallback
+    datasets: ${eval_dataset}
+
+logger:
+  wandb:
+    group: sample_enhanced_sampling_data
+    notes: "Sampling from enhanced sampling data using memory sampler"
+    tags: ["sample", "enhanced_sampling", "memory_sampler", "ALA_ALA", "spatiotemporal conditioner"]