Skills › Scientific Agent Skills › diffdock › prepare_batch_csv.py
prepare_batch_csv.py
289 lines · skills/diffdock/scripts/prepare_batch_csv.py at 5f2bc03
1#!/usr/bin/env python3
2"""
3DiffDock Batch CSV Preparation and Validation Script
5This script helps prepare and validate CSV files for DiffDock batch processing.
6It checks for required columns, validates file paths, and ensures SMILES strings
7are properly formatted.
9Usage:
10 python prepare_batch_csv.py input.csv --validate
11 python prepare_batch_csv.py --create --output batch_input.csv
12"""
14import argparse
15import csv
16import re
17import os
18import sys
19import pandas as pd
20from pathlib import Path
22try:
23 from rdkit import Chem
24 from rdkit import RDLogger
25 RDLogger.DisableLog('rdApp.*')
26 RDKIT_AVAILABLE = True
27except ImportError:
28 RDKIT_AVAILABLE = False
29 print("[WARN] RDKit not available. SMILES-only rows cannot pass validation.")
32def validate_smiles(smiles_string):
33 """Validate a SMILES string using RDKit."""
34 if not RDKIT_AVAILABLE:
35 return False, "RDKit not available; chemistry validation was not performed"
37 try:
38 mol = Chem.MolFromSmiles(smiles_string)
39 if mol is None or mol.GetNumAtoms() == 0:
40 return False, "Invalid SMILES structure"
41 return True, "Valid SMILES"
42 except Exception as e:
43 return False, str(e)
46def validate_file_path(file_path, base_dir=None):
47 """Validate that a file path exists."""
48 if pd.isna(file_path) or file_path == "":
49 return True, "Empty (will use protein_sequence)"
51 # Handle relative paths
52 if base_dir:
53 full_path = Path(base_dir) / file_path
54 else:
55 full_path = Path(file_path)
57 if full_path.is_file():
58 return True, f"File exists: {full_path}"
59 else:
60 return False, f"File not found: {full_path}"
63def validate_csv(csv_path, base_dir=None):
64 """
65 Validate a DiffDock batch input CSV file.
67 Args:
68 csv_path: Path to CSV file
69 base_dir: Inference working directory for relative paths (default: current directory)
71 Returns:
72 bool: True if validation passes
73 list: List of validation messages
74 """
75 messages = []
76 valid = True
78 # Read CSV
79 try:
80 with open(csv_path, newline="", encoding="utf-8-sig") as handle:
81 reader = csv.reader(handle, strict=True)
82 header = next(reader)
83 if len(header) != len(set(header)):
84 raise ValueError("Duplicate CSV column names")
85 for line_number, fields in enumerate(reader, 2):
86 if len(fields) != len(header):
87 raise ValueError(f"CSV row {line_number} has {len(fields)} fields; expected {len(header)}")
88 df = pd.read_csv(csv_path, dtype=str, keep_default_na=False, encoding="utf-8-sig")
89 messages.append(f"[OK] Successfully read CSV with {len(df)} rows")
90 except Exception as e:
91 messages.append(f"[FAIL] Error reading CSV: {e}")
92 return False, messages
94 # Check required columns
95 required_cols = ['complex_name', 'protein_path', 'ligand_description', 'protein_sequence']
96 missing_cols = [col for col in required_cols if col not in df.columns]
98 if missing_cols:
99 messages.append(f"[FAIL] Missing required columns: {', '.join(missing_cols)}")
100 valid = False
101 else:
102 messages.append("[OK] All required columns present")
104 # Set base directory
105 if base_dir is None:
106 base_dir = Path.cwd()
108 if df.empty:
109 messages.append("[FAIL] CSV contains no complexes")
110 valid = False
111 seen_names = set()
113 # Validate each row. The per-row checks index the required columns
114 # directly, so they can only run once every one of them is present --
115 # otherwise a CSV missing a column raises KeyError instead of reporting it.
116 rows = [] if missing_cols else df.iterrows()
117 for idx, row in rows:
118 row_msgs = []
120 # Names become output directory names and ESM labels upstream.
121 name = row['complex_name']
122 if not name:
123 row_msgs.append("Missing complex_name")
124 valid = False
125 elif not re.fullmatch(r"[A-Za-z_][A-Za-z0-9_.-]*", name) or name in {"NA", "N/A", "NULL", "null", "None", "NaN", "nan"}:
126 row_msgs.append("Unsafe complex_name; use a letter/underscore followed by letters, digits, _, . or - (not an NA token)")
127 valid = False
128 if name in seen_names:
129 row_msgs.append("Duplicate complex_name would overwrite output")
130 valid = False
131 seen_names.add(name)
133 # Check that either protein_path or protein_sequence is provided
134 has_protein_path = not pd.isna(row['protein_path']) and row['protein_path'] != ""
135 has_protein_seq = not pd.isna(row['protein_sequence']) and row['protein_sequence'] != ""
137 if not has_protein_path and not has_protein_seq:
138 row_msgs.append("Must provide either protein_path or protein_sequence")
139 valid = False
140 elif has_protein_path and has_protein_seq:
141 row_msgs.append("Warning: Both protein_path and protein_sequence provided, will use protein_path")
143 # Validate protein path if provided
144 if has_protein_path:
145 file_valid, msg = validate_file_path(row['protein_path'], base_dir)
146 if not file_valid:
147 row_msgs.append(f"Protein file issue: {msg}")
148 valid = False
150 if has_protein_seq and not has_protein_path:
151 sequence = row['protein_sequence']
152 if not re.fullmatch(r"[ACDEFGHIKLMNPQRSTVWYXBZUO]+(?::[ACDEFGHIKLMNPQRSTVWYXBZUO]+)*", sequence):
153 row_msgs.append("Invalid protein_sequence; provide full uppercase amino-acid sequence, no ellipsis")
154 valid = False
155 elif any(len(chain) > 1022 for chain in sequence.split(':')):
156 row_msgs.append("Protein chain exceeds upstream inference ESM truncation length (1022 residues)")
157 valid = False
159 # Upstream tries SMILES first: / and backslash encode stereochemistry.
160 ligand_desc = row['ligand_description']
161 if not ligand_desc:
162 row_msgs.append("Missing ligand_description")
163 valid = False
164 else:
165 smiles_valid, msg = validate_smiles(ligand_desc)
166 if not smiles_valid:
167 suffix = Path(ligand_desc).suffix
168 if suffix in {'.sdf', '.mol2', '.pdb', '.pdbqt'}:
169 file_valid, file_msg = validate_file_path(ligand_desc, base_dir)
170 if not file_valid:
171 row_msgs.append(f"Ligand file issue: {file_msg}")
172 valid = False
173 else:
174 row_msgs.append(f"SMILES issue: {msg}")
175 valid = False
177 if row_msgs:
178 messages.append(f"\nRow {idx + 1} ({row.get('complex_name', 'unnamed')}):")
179 for msg in row_msgs:
180 messages.append(f" - {msg}")
182 # Summary
183 messages.append(f"\n{'='*60}")
184 if valid:
185 messages.append("[OK] CSV validation PASSED - input paths and SMILES checked; docking and file chemistry not validated")
186 else:
187 messages.append("[FAIL] CSV validation FAILED - please fix issues above")
189 return valid, messages
192def create_template_csv(output_path, num_examples=3):
193 """Create a template CSV file with example entries."""
195 if num_examples < 1:
196 raise ValueError("num_examples must be at least 1")
197 examples = {
198 'complex_name': ['example1', 'example2', 'example3'][:num_examples],
199 'protein_path': ['protein1.pdb', '', 'protein3.pdb'][:num_examples],
200 'ligand_description': [
201 'CC(=O)Oc1ccccc1C(=O)O', # Aspirin SMILES
202 'COc1ccc(C#N)cc1', # Example SMILES
203 'ligand.sdf' # Example file path
204 ][:num_examples],
205 'protein_sequence': [
206 '', # Empty - using PDB file
207 'MSKGEELFTGVVPILVELDGDVNGHKFSVSGEGEGDATYGKLTLKFICTTGKLPVPWPTLVTTFSYGVQCFSRYPDHMKQHDFFKSAMPEGYVQERTIFFKDDGNYKTRAEVKFEGDTLVNRIELKGIDFKEDGNILGHKLEYNYNSHNVYIMADKQKNGIKVNFKIRHNIEDGSVQLADHYQQNTPIGDGPVLLPDNHYLSTQSALSKDPNEKRDHMVLLEFVTAAGITHGMDELYK', # GFP sequence
208 '' # Empty - using PDB file
209 ][:num_examples]
210 }
212 df = pd.DataFrame(examples)
213 df.to_csv(output_path, index=False)
215 return df
218def main():
219 parser = argparse.ArgumentParser(
220 description='Prepare and validate DiffDock batch CSV files',
221 formatter_class=argparse.RawDescriptionHelpFormatter,
222 epilog="""
223Examples:
224 # Validate existing CSV
225 python prepare_batch_csv.py input.csv --validate
227 # Create template CSV
228 python prepare_batch_csv.py --create --output batch_template.csv
230 # Create template with 2 example rows
231 python prepare_batch_csv.py --create --output template.csv --num-examples 2
233 # Validate with custom base directory for relative paths
234 python prepare_batch_csv.py input.csv --validate --base-dir /path/to/data/
235 """
236 )
238 parser.add_argument('csv_file', nargs='?', help='CSV file to validate')
239 parser.add_argument('--validate', action='store_true',
240 help='Validate the CSV file')
241 parser.add_argument('--create', action='store_true',
242 help='Create a template CSV file')
243 parser.add_argument('--output', '-o', help='Output path for template CSV')
244 parser.add_argument('--num-examples', type=int, default=3,
245 help='Number of example rows in template, 1-3 (default: 3)')
246 parser.add_argument('--base-dir', help='Planned inference working directory for relative paths (default: current directory)')
248 args = parser.parse_args()
250 # Create template
251 if args.create:
252 output_path = args.output or 'diffdock_batch_template.csv'
253 df = create_template_csv(output_path, args.num_examples)
254 print(f"[OK] Created template CSV: {output_path}")
255 print(f"\nTemplate contents:")
256 print(df.to_string(index=False))
257 print(f"\nEdit this file with your protein-ligand pairs and run with:")
258 print(f" python -m inference --config default_inference_args.yaml \\")
259 print(f" --protein_ligand_csv {output_path} --out_dir results/")
260 return 0
262 # Validate CSV
263 if args.validate or args.csv_file:
264 if not args.csv_file:
265 print("Error: CSV file required for validation")
266 parser.print_help()
267 return 1
269 if not os.path.exists(args.csv_file):
270 print(f"Error: CSV file not found: {args.csv_file}")
271 return 1
273 print(f"Validating: {args.csv_file}")
274 print("="*60)
276 valid, messages = validate_csv(args.csv_file, args.base_dir)
278 for msg in messages:
279 print(msg)
281 return 0 if valid else 1
283 # No action specified
284 parser.print_help()
285 return 1
288if __name__ == '__main__':
289 sys.exit(main())