聚类方法,聚类后选择打分最高那个分子,并对 karamadock 的结果求交集

scripts/extract_top_molecules.py

@@ -0,0 +1,64 @@
+import pandas as pd
+import os
+import ast
+import argparse
+
+def parse_vina_scores(vina_scores_str):
+    """解析vina_scores字符串为浮点数列表"""
+    try:
+        scores = ast.literal_eval(vina_scores_str)
+        if isinstance(scores, list) and len(scores) > 0:
+            return scores[0]  # 取第一个值作为vina_score
+        return None
+    except:
+        return None
+
+def extract_top_molecules(file_path, output_dir, dataset_name):
+    """从CSV文件中提取karma_score_aligned和vina_score前1000的分子"""
+    # 读取数据
+    df = pd.read_csv(file_path)
+    
+    # 解析vina_scores列
+    df['vina_score'] = df['vina_scores'].apply(parse_vina_scores)
+    
+    # 按karma_score_aligned排序并提取前1000
+    df_karma_top = df.sort_values('karma_score_aligned', ascending=False).head(1000)
+    
+    # 按vina_score排序并提取前1000
+    df_vina_top = df.sort_values('vina_score', ascending=False).head(1000)
+    
+    # 保存结果
+    karma_output_file = os.path.join(output_dir, f"{dataset_name}_karma_score_aligned_top1000.csv")
+    vina_output_file = os.path.join(output_dir, f"{dataset_name}_vina_score_top1000.csv")
+    
+    df_karma_top.to_csv(karma_output_file, index=False)
+    df_vina_top.to_csv(vina_output_file, index=False)
+    
+    print(f"{dataset_name} - karma_score_aligned前1000分子保存到: {karma_output_file}")
+    print(f"{dataset_name} - vina_score前1000分子保存到: {vina_output_file}")
+    print(f"{dataset_name} - karma_score_aligned前1000分子数量: {len(df_karma_top)}")
+    print(f"{dataset_name} - vina_score前1000分子数量: {len(df_vina_top)}")
+    
+    return df_karma_top, df_vina_top
+
+def main():
+    parser = argparse.ArgumentParser(description='从CSV文件中提取karma_score_aligned和vina_score前1000的分子')
+    parser.add_argument('--input', nargs='+', required=True, 
+                        help='输入CSV文件路径列表')
+    parser.add_argument('--dataset-names', nargs='+', required=True,
+                        help='数据集名称列表，与输入文件一一对应')
+    parser.add_argument('--output', required=True,
+                        help='输出目录')
+    
+    args = parser.parse_args()
+    
+    # 确保输出目录存在
+    os.makedirs(args.output, exist_ok=True)
+    
+    # 处理每个文件
+    for file_path, dataset_name in zip(args.input, args.dataset_names):
+        print(f"Processing {dataset_name}...")
+        extract_top_molecules(file_path, args.output, dataset_name)
+
+if __name__ == "__main__":
+    main()