Skills › Scientific Agent Skills › diffdock › analyze_results.py
analyze_results.py
357 lines · skills/diffdock/scripts/analyze_results.py at 5f2bc03
1#!/usr/bin/env python3
2"""
3DiffDock Results Analysis Script
5This script analyzes DiffDock prediction results, extracting confidence scores,
6ranking predictions, and generating summary reports.
8Usage:
9 python analyze_results.py results/output_dir/
10 python analyze_results.py results/ --top 50 --threshold 0.0
11 python analyze_results.py results/ --export summary.csv
12"""
14import argparse
15import os
16import sys
17import math
18from pathlib import Path
19from collections import defaultdict
20import re
23def parse_confidence_scores(results_dir):
24 """
25 Parse confidence scores from DiffDock output directory.
27 Args:
28 results_dir: Path to DiffDock results directory
30 Returns:
31 dict: Dictionary mapping complex names to their predictions and scores
32 """
33 results = {}
34 results_path = Path(results_dir)
36 # An unrelated input SDF in a batch root must not hide all nested poses.
37 parsed = parse_single_complex(results_path)
38 for subdir in sorted(results_path.iterdir()):
39 if subdir.is_dir():
40 complex_results = parse_single_complex(subdir)
41 if complex_results:
42 results[subdir.name] = complex_results
43 if parsed:
44 if results:
45 raise ValueError('Mixed single-complex and batch pose layout; select one run directory')
46 results['single_complex'] = parsed
48 return results
51def parse_single_complex(complex_dir):
52 """Parse results for a single complex."""
53 predictions_by_rank = {}
55 # Look for SDF files with rank information
56 for sdf_file in complex_dir.glob("*.sdf"):
57 filename = sdf_file.name
59 # Current DiffDock writes rank1_confidence0.87.sdf; older runs may use rank_1.sdf.
60 rank_match = re.match(r'^rank_?([1-9]\d*)(?:_|\.)', filename)
61 if rank_match:
62 rank = int(rank_match.group(1))
64 # Try to extract confidence score from filename or separate file
65 confidence = extract_confidence_score(sdf_file, complex_dir)
67 prediction = {
68 'rank': rank,
69 'file': sdf_file.name,
70 'path': str(sdf_file),
71 'confidence': confidence
72 }
74 existing = predictions_by_rank.get(rank)
75 if existing and 'confidence' in filename and 'confidence' in existing['file']:
76 raise ValueError(f'Duplicate scored files for rank {rank} in {complex_dir}; use a fresh output directory')
77 if (
78 existing is None
79 or (existing['confidence'] is None and confidence is not None)
80 or ('confidence' in filename and 'confidence' not in existing['file'])
81 ):
82 predictions_by_rank[rank] = prediction
84 # Sort by rank
85 predictions = list(predictions_by_rank.values())
86 predictions.sort(key=lambda x: x['rank'])
88 return {'predictions': predictions} if predictions else None
91def extract_confidence_score(sdf_file, complex_dir):
92 """
93 Extract confidence score for a prediction.
95 Tries multiple methods:
96 1. Extract from current filename format, e.g. rank1_confidence0.87.sdf
97 2. Read from legacy confidence_scores.txt file
98 3. Parse from SDF file properties
99 """
100 # Method 1: Filename
101 conf_match = re.search(r'(?:confidence|conf)_?([-+]?(?:\d+(?:\.\d*)?|\.\d+)(?:[eE][-+]?\d+)?)', sdf_file.name)
102 if conf_match:
103 return finite_score(conf_match.group(1))
105 # Method 2: legacy confidence_scores.txt
106 confidence_file = complex_dir / "confidence_scores.txt"
107 if confidence_file.exists():
108 try:
109 with open(confidence_file) as f:
110 lines = f.readlines()
111 # Extract rank from filename
112 rank_match = re.search(r'rank_?(\d+)', sdf_file.name)
113 if rank_match:
114 rank = int(rank_match.group(1))
115 if 1 <= rank <= len(lines):
116 return finite_score(lines[rank - 1].strip())
117 except Exception:
118 pass
120 # Method 3: Parse from SDF file
121 try:
122 with open(sdf_file) as f:
123 content = f.read()
124 # Look for confidence score in SDF properties. An SDF data item is
125 # written as `> <confidence>` followed by the value on the next
126 # line, so the separator class has to admit the closing angle
127 # bracket as well as a colon.
128 conf_match = re.search(r'confidence[>:\s]+([-+]?(?:\d+(?:\.\d*)?|\.\d+)(?:[eE][-+]?\d+)?)', content, re.IGNORECASE)
129 if conf_match:
130 return finite_score(conf_match.group(1))
131 except Exception:
132 pass
134 return None
137def finite_score(value):
138 """Nonfinite scores carry no interpretable confidence."""
139 score = float(value)
140 return score if math.isfinite(score) else None
143def classify_confidence(score):
144 """Classify confidence score into categories."""
145 if score is None or not math.isfinite(score):
146 return "Unknown"
147 elif score > 0:
148 return "High"
149 elif score > -1.5:
150 return "Moderate"
151 else:
152 return "Low"
155def print_summary(results, top_n=None, min_confidence=None):
156 """Print a formatted summary of results."""
158 print("\n" + "="*80)
159 print("DiffDock Results Summary")
160 print("="*80)
162 all_predictions = []
164 for complex_name, data in results.items():
165 predictions = data.get('predictions', [])
167 print(f"\n{complex_name}")
168 print("-" * 80)
170 if not predictions:
171 print(" No predictions found")
172 continue
174 # Filter by confidence if specified
175 filtered_predictions = predictions
176 if min_confidence is not None:
177 filtered_predictions = [p for p in predictions if p['confidence'] is not None and p['confidence'] >= min_confidence]
179 # Limit to top N if specified
180 if top_n is not None:
181 filtered_predictions = filtered_predictions[:top_n]
183 for pred in filtered_predictions:
184 confidence = pred['confidence']
185 confidence_class = classify_confidence(confidence)
187 conf_str = f"{confidence:>7.3f}" if confidence is not None else " N/A"
188 print(f" Rank {pred['rank']:2d}: Confidence = {conf_str} ({confidence_class:8s}) | {pred['file']}")
190 # Add to all predictions for overall statistics
191 if confidence is not None:
192 all_predictions.append((complex_name, pred['rank'], confidence))
194 # Show statistics for this complex
195 if filtered_predictions and any(p['confidence'] is not None for p in filtered_predictions):
196 confidences = [p['confidence'] for p in filtered_predictions if p['confidence'] is not None]
197 print(f"\n Statistics: {len(filtered_predictions)} predictions")
198 print(f" Mean confidence: {sum(confidences)/len(confidences):.3f}")
199 print(f" Max confidence: {max(confidences):.3f}")
200 print(f" Min confidence: {min(confidences):.3f}")
202 # Overall statistics
203 if all_predictions:
204 print("\n" + "="*80)
205 print("Overall Statistics")
206 print("="*80)
208 confidences = [conf for _, _, conf in all_predictions]
209 print(f" Total predictions: {len(all_predictions)}")
210 print(f" Total complexes: {len(results)}")
211 print(f" Mean confidence: {sum(confidences)/len(confidences):.3f}")
212 print(f" Max confidence: {max(confidences):.3f}")
213 print(f" Min confidence: {min(confidences):.3f}")
215 # Confidence distribution
216 high = sum(1 for c in confidences if c > 0)
217 moderate = sum(1 for c in confidences if -1.5 < c <= 0)
218 low = sum(1 for c in confidences if c <= -1.5)
220 print(f"\n Confidence distribution:")
221 print(f" High (> 0): {high:4d} ({100*high/len(confidences):5.1f}%)")
222 print(f" Moderate (-1.5 to 0): {moderate:4d} ({100*moderate/len(confidences):5.1f}%)")
223 print(f" Low (<= -1.5): {low:4d} ({100*low/len(confidences):5.1f}%)")
225 print("\n" + "="*80)
228def export_to_csv(results, output_path):
229 """Export results to CSV file."""
230 import csv
232 with open(output_path, 'w', newline='') as f:
233 writer = csv.writer(f)
234 writer.writerow(['complex_name', 'rank', 'confidence', 'confidence_class', 'file_path'])
236 for complex_name, data in results.items():
237 predictions = data.get('predictions', [])
238 for pred in predictions:
239 confidence = pred['confidence']
240 confidence_class = classify_confidence(confidence)
241 conf_value = confidence if confidence is not None else ''
243 writer.writerow([
244 complex_name,
245 pred['rank'],
246 conf_value,
247 confidence_class,
248 pred['path']
249 ])
251 print(f"[OK] Exported results to: {output_path}")
254def get_top_predictions(results, n=10, sort_by='confidence'):
255 """Get top N predictions across all complexes."""
256 all_predictions = []
258 for complex_name, data in results.items():
259 predictions = data.get('predictions', [])
260 for pred in predictions:
261 if pred['confidence'] is not None:
262 all_predictions.append({
263 'complex': complex_name,
264 **pred
265 })
267 # Sort by confidence (descending)
268 all_predictions.sort(key=lambda x: x['confidence'], reverse=True)
270 return all_predictions[:n]
273def print_top_predictions(results, n=10):
274 """Print top N predictions across all complexes."""
275 top_preds = get_top_predictions(results, n)
277 print("\n" + "="*80)
278 print(f"Top {n} Predictions Across All Complexes")
279 print("Confidence triage only: scores across complexes are not calibrated affinity rankings.")
280 print("="*80)
282 for i, pred in enumerate(top_preds, 1):
283 confidence_class = classify_confidence(pred['confidence'])
284 print(f"{i:2d}. {pred['complex']:30s} | Rank {pred['rank']:2d} | "
285 f"Confidence: {pred['confidence']:7.3f} ({confidence_class})")
287 print("="*80)
290def main():
291 parser = argparse.ArgumentParser(
292 description='Analyze DiffDock prediction results',
293 formatter_class=argparse.RawDescriptionHelpFormatter,
294 epilog="""
295Examples:
296 # Analyze all results in directory
297 python analyze_results.py results/output_dir/
299 # Show only top 5 predictions per complex
300 python analyze_results.py results/ --top 5
302 # Filter by confidence threshold
303 python analyze_results.py results/ --threshold 0.0
305 # Export to CSV
306 python analyze_results.py results/ --export summary.csv
308 # Show top 20 predictions across all complexes
309 python analyze_results.py results/ --best 20
310 """
311 )
313 parser.add_argument('results_dir', help='Path to DiffDock results directory')
314 parser.add_argument('--top', '-t', type=int,
315 help='Show only top N predictions per complex')
316 parser.add_argument('--threshold', type=float,
317 help='Minimum confidence threshold')
318 parser.add_argument('--export', '-e', metavar='FILE',
319 help='Export results to CSV file')
320 parser.add_argument('--best', '-b', type=int, metavar='N',
321 help='Show top N predictions across all complexes')
323 args = parser.parse_args()
325 # Validate results directory
326 if not Path(args.results_dir).is_dir():
327 print(f"Error: Results directory not found: {args.results_dir}")
328 return 1
330 # Parse results
331 print(f"Analyzing results in: {args.results_dir}")
332 try:
333 results = parse_confidence_scores(args.results_dir)
334 except ValueError as error:
335 print(f"Error: {error}")
336 return 1
338 if not results:
339 print("No DiffDock results found in directory")
340 return 1
342 # Print summary
343 print_summary(results, top_n=args.top, min_confidence=args.threshold)
345 # Print top predictions across all complexes
346 if args.best:
347 print_top_predictions(results, args.best)
349 # Export to CSV if requested
350 if args.export:
351 export_to_csv(results, args.export)
353 return 0
356if __name__ == '__main__':
357 sys.exit(main())