代码移动位置

scripts/analyze_qed_mw_distribution.py

@@ -17,11 +17,6 @@ import logging
 import ast
 import json
 import click
-import sys
-import os
-
-# Add the parent directory to the path to import modules
-sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
 
 # Setup logging
 logging.basicConfig(level=logging.INFO)
@@ -125,8 +120,8 @@ def load_vina_scores_from_csv(df, max_files=1000):
         if processed_files >= max_files:
             break
             
-        # Skip reference molecules (those with align_ and _out_converted.sdf in their filename)
-        if 'align_' in row['filename'] and '_out_converted.sdf' in row['filename']:
+        # Skip reference molecules (those with mol2 extension)
+        if '.mol2' in row['filename']:
             continue
             
         try:
@@ -179,8 +174,8 @@ def get_reference_vina_scores(dataset_name, rank=0):
     """
     reference_scores = {}
     
-    # 使用更新后的路径以适应新目录结构
-    reference_dir = Path("../result") / "refence" / dataset_name
+    # 使用原始目录名称 "refence"
+    reference_dir = Path("result") / "refence" / dataset_name
     
     if not reference_dir.exists():
         logger.warning(f"Reference directory {reference_dir} does not exist")
@@ -207,7 +202,7 @@ def get_reference_vina_scores(dataset_name, rank=0):
             if '_addH' in filename_stem:
                 filename_stem = filename_stem.replace('_addH', '')
             if 'align_' in filename_stem:
-                filename_stem = filename_stem.split('_')[-1]
+                filename_stem = filename_stem.split('_')[-1]  # Get the last part (e.g., 9NY or 0GA)
             
             # Use filename_stem as key for reference_scores
             reference_scores[filename_stem] = reference_score