Random sample

input_csv_split/split_csv.sh

@@ -0,0 +1,16 @@
+#!/bin/bash
+# Usage: ./split_csv.sh input.csv output_dir
+
+input_csv="$1"
+output_dir="$2"
+
+mkdir -p "$output_dir"
+
+# Read header
+header=$(head -n 1 "$input_csv")
+
+tail -n +2 "$input_csv" | while IFS=, read -r name seqres; do
+    name_clean=$(echo "$name" | tr -d '\r\n ')
+    echo "$header" > "$output_dir/${name_clean}.csv"
+    echo "$name,$seqres" >> "$output_dir/${name_clean}.csv"
+done

pyproject.toml

@@ -52,6 +52,7 @@ training = ["wandb==0.16.6",
             "openfold @ git+https://github.com/aqlaboratory/openfold.git@103d037",
             "ml-collections==0.1.0",
             "dm-tree==0.1.6",
+            "lightning==2.1.0",
             "modelcif==0.7",
             "biopython==1.79",   # specific versions of main dependencies for reproducibility
             "pandas==1.5.3",

scripts/run_ensemble_model/DEBUG_README.md

@@ -0,0 +1,137 @@
+# Debug Scripts for train_structure.py
+
+After removing redundant files (`new_train.py`, `simple_test.py`), use these scripts to verify everything still works.
+
+## Quick Import Test (Fastest - 10 seconds)
+
+Test all Python imports without running actual training:
+
+```bash
+cd /pscratch/sd/l/lemonboy/metfish/scripts/run_ensemble_model
+python test_imports.py
+```
+
+This checks:
+- All required imports work
+- `StructureModel` can be instantiated
+- Methods like `get_trainable_parameters()` exist
+
+**Expected output:** `SUCCESS: All tests passed!`
+
+---
+
+## Full Debug Run (Medium - 5-10 minutes)
+
+Run a quick 50-iteration optimization on one sequence:
+
+```bash
+cd /pscratch/sd/l/lemonboy/metfish/scripts/run_ensemble_model
+
+# Option 1: Use default test protein (1AEL-12_A)
+./debug_train_structure.sh
+
+# Option 2: Test specific protein
+./debug_train_structure.sh /path/to/your/input.csv
+```
+
+**Output location:** `/pscratch/sd/l/lemonboy/metfish/debug_output/train_structure_test_*`
+
+**What it does:**
+- Runs 50 iterations (instead of 500) for quick testing
+- Saves to separate debug directory (won't overwrite production data)
+- Creates timestamped output folder
+- Verifies PDB files are generated
+
+**Expected files:**
+```
+debug_output/train_structure_test_PROTEIN_20251209_HHMMSS/
+├── initial/
+│   └── PROTEIN.pdb
+├── intermediate/
+│   ├── PROTEIN_iter_0010.pdb
+│   ├── PROTEIN_iter_0020.pdb
+│   └── ...
+├── best/
+│   └── PROTEIN_iter_XXXX.pdb
+├── final/
+│   └── PROTEIN_optimized.pdb
+├── best_model.pth
+├── loss_history.npy
+└── optimization_summary.txt
+```
+
+---
+
+## Production Run (Full - Hours)
+
+Once debug tests pass, run full 500-iteration optimization:
+
+```bash
+# Submit to SLURM
+cd /pscratch/sd/l/lemonboy/metfish/scripts/run_ensemble_model
+sbatch train_structure.slurm /path/to/input.csv
+```
+
+---
+
+## Troubleshooting
+
+### Import test fails
+```bash
+# Check conda environment
+conda activate /pscratch/sd/l/lemonboy/AlphaSAXS_2025
+which python
+
+# Try importing manually
+python -c "from metfish.refinement_model.random_model import StructureModel; print('OK')"
+```
+
+### Debug run fails
+Check error messages in:
+- Terminal output
+- SLURM output file (if using sbatch)
+
+Common issues:
+- **Missing checkpoint:** Verify `CKPT_PATH` exists
+- **Missing data:** Verify `DATA_DIR` and input CSV exist
+- **CUDA errors:** Ensure GPU is available (for SLURM runs)
+
+### Files not generated
+Check:
+```bash
+ls -lh /pscratch/sd/l/lemonboy/metfish/debug_output/
+```
+
+If empty, check permissions and disk space:
+```bash
+df -h /pscratch/sd/l/lemonboy/metfish/
+```
+
+---
+
+## Key Changes After Cleanup
+
+✅ **Removed:**
+- `new_train.py` (unused for paper data generation)
+- `simple_test.py` (empty file)
+
+✅ **Kept:**
+- `train_structure.py` - main optimization script (YOUR PAPER DATA)
+- `save_structure_output()` in `train_structure.py` (as requested)
+- `MSARandomModel` in `random_model.py` (used by `generate.py`)
+- `compute_plddt_loss()` in `random_model.py` (imported by `train_structure.py`)
+
+✅ **Your production command still works:**
+```bash
+python $SCRATCH/metfish/src/metfish/refinement_model/train_structure.py \
+    --data_dir "$DATA_DIR" \
+    --output_dir "$OUTPUT_DIR" \
+    --ckpt_path "$CKPT_PATH" \
+    --test_csv_name "$INPUT_CSV" \
+    --saxs_ext _atom_only.csv \
+    --num_iterations 500 \
+    --learning_rate 1e-3 \
+    --sequence_index 0 \
+    --save_frequency 5 \
+    --random_init
+```

scripts/run_ensemble_model/debug_train_structure.sh

@@ -0,0 +1,102 @@
+#!/bin/bash
+# Debug script to test train_structure.py on a single sequence
+# Usage: bash debug_train_structure.sh [input_csv]
+
+set -e  # Exit on error
+
+# Activate conda environment
+echo "Activating conda environment..."
+source ~/.bashrc
+conda activate /pscratch/sd/l/lemonboy/AlphaSAXS_2025
+
+# Set directories
+DATA_DIR=/global/cfs/cdirs/m4704/100125_Nature_Com_data/Apo_holo_data
+CKPT_PATH=/global/cfs/cdirs/m4704/100125_Nature_Com_data/ensemble_generated/checkpoint/epoch=15-step=21009.ckpt
+
+# Input CSV - use provided argument or default test file
+if [ -z "$1" ]; then
+    echo "No input CSV provided, using default test file..."
+    INPUT_CSV="$DATA_DIR/input_csv_split/1AEL-12_A.csv"  # Use first protein as test
+else
+    INPUT_CSV="$1"
+fi
+
+# Resolve to absolute path
+if [[ "$INPUT_CSV" != /* ]]; then
+    INPUT_CSV="$(realpath "$INPUT_CSV")"
+fi
+
+CSV_BASENAME=$(basename "$INPUT_CSV" .csv)
+
+# DEBUG OUTPUT DIR - separate location to avoid overwriting production data
+OUTPUT_DIR=/pscratch/sd/l/lemonboy/metfish/debug_output/train_structure_test_${CSV_BASENAME}_$(date +%Y%m%d_%H%M%S)
+mkdir -p "$OUTPUT_DIR"
+
+echo "=========================================="
+echo "DEBUG TEST: train_structure.py"
+echo "=========================================="
+echo "Data directory: $DATA_DIR"
+echo "Checkpoint: $CKPT_PATH"
+echo "Input CSV: $INPUT_CSV"
+echo "Output directory: $OUTPUT_DIR"
+echo "=========================================="
+
+# Verify files exist
+if [ ! -f "$INPUT_CSV" ]; then
+    echo "ERROR: Input CSV not found: $INPUT_CSV"
+    exit 1
+fi
+
+if [ ! -f "$CKPT_PATH" ]; then
+    echo "ERROR: Checkpoint not found: $CKPT_PATH"
+    exit 1
+fi
+
+# Run with reduced iterations for quick test
+echo "Starting optimization test (50 iterations for quick debug)..."
+python $SCRATCH/metfish/src/metfish/refinement_model/train_structure.py \
+    --data_dir "$DATA_DIR" \
+    --output_dir "$OUTPUT_DIR" \
+    --ckpt_path "$CKPT_PATH" \
+    --test_csv_name "$INPUT_CSV" \
+    --saxs_ext _atom_only.csv \
+    --num_iterations 50 \
+    --learning_rate 1e-3 \
+    --sequence_index 0 \
+    --save_frequency 10 \
+    --random_init
+
+# Check if output files were created
+echo ""
+echo "=========================================="
+echo "DEBUG TEST RESULTS"
+echo "=========================================="
+if [ -d "$OUTPUT_DIR" ]; then
+    echo "✓ Output directory created: $OUTPUT_DIR"
+    echo ""
+    echo "Files generated:"
+    find "$OUTPUT_DIR" -type f -name "*.pdb" -o -name "*.txt" -o -name "*.npy" -o -name "*.pth" | head -20
+    echo ""
+    echo "Directory structure:"
+    tree -L 2 "$OUTPUT_DIR" 2>/dev/null || ls -R "$OUTPUT_DIR"
+    echo ""
+    echo "✓ Test completed successfully!"
+    echo "Review results in: $OUTPUT_DIR"
+else
+    echo "✗ ERROR: Output directory not created"
+    exit 1
+fi
+
+echo ""
+echo "To run full optimization (500 iterations), use:"
+echo "python \$SCRATCH/metfish/src/metfish/refinement_model/train_structure.py \\"
+echo "    --data_dir \"$DATA_DIR\" \\"
+echo "    --output_dir \"/path/to/full/output\" \\"
+echo "    --ckpt_path \"$CKPT_PATH\" \\"
+echo "    --test_csv_name \"$INPUT_CSV\" \\"
+echo "    --saxs_ext _atom_only.csv \\"
+echo "    --num_iterations 500 \\"
+echo "    --learning_rate 1e-3 \\"
+echo "    --sequence_index 0 \\"
+echo "    --save_frequency 5 \\"
+echo "    --random_init"