SIME/utils/mole_predictor.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
MolE 抗菌活性预测工具

这个脚本提供了使用 MolE 模型预测小分子 SMILES 抗菌活性的功能。
支持命令行和 Python API 调用两种方式。

命令行示例:
    python mole_predictor.py input.csv output.csv --smiles-column smiles --id-column chem_id
    
Python API 示例:
    from utils.mole_predictor import predict_csv_file
    
    predict_csv_file(
        input_path="input.csv",
        output_path="output.csv",
        smiles_column="smiles",
        id_column="chem_id"
    )
"""

import sys
import os
from pathlib import Path

# 添加项目根目录到 Python 路径
project_root = Path(__file__).parent.parent
sys.path.insert(0, str(project_root))

import click
import pandas as pd
from typing import Optional, List
from datetime import datetime
from tqdm import tqdm

from models.broad_spectrum_predictor import (
    ParallelBroadSpectrumPredictor,
    PredictionConfig,
    MoleculeInput,
    BroadSpectrumResult
)


def predict_csv_file(
    input_path: str,
    output_path: Optional[str] = None,
    smiles_column: str = "smiles",
    id_column: str = "chem_id",
    batch_size: int = 100,
    n_workers: Optional[int] = None,
    device: str = "auto",
    add_suffix: bool = True,
    include_strain_predictions: bool = False
) -> pd.DataFrame:
    """
    预测 CSV 文件中的分子抗菌活性
    
    Args:
        input_path: 输入 CSV 文件路径
        output_path: 输出 CSV 文件路径，如果为 None 则自动生成
        smiles_column: SMILES 列名
        id_column: 化合物 ID 列名
        batch_size: 批处理大小
        n_workers: 工作进程数
        device: 计算设备 ("auto", "cpu", "cuda:0" 等)
        add_suffix: 是否在输出文件名后添加预测后缀
        include_strain_predictions: 是否在输出中包含40种菌株的预测详情
    
    Returns:
        包含预测结果的 DataFrame
    """
    
    print(f"开始处理文件: {input_path}")
    
    # 读取输入文件
    input_path_obj = Path(input_path)
    if not input_path_obj.exists():
        raise FileNotFoundError(f"输入文件不存在: {input_path}")
    
    # 读取 CSV
    try:
        df_input = pd.read_csv(input_path)
    except Exception as e:
        raise RuntimeError(f"读取 CSV 文件失败: {e}")
    
    print(f"读取了 {len(df_input)} 条数据")
    
    # 检查列是否存在（大小写不敏感）
    columns_lower = {col.lower(): col for col in df_input.columns}
    
    smiles_col_actual = columns_lower.get(smiles_column.lower())
    if smiles_col_actual is None:
        raise ValueError(
            f"未找到 SMILES 列 '{smiles_column}'。可用列: {list(df_input.columns)}"
        )
    
    # 处理 ID 列
    id_col_actual = columns_lower.get(id_column.lower())
    if id_col_actual is None:
        print(f"未找到 ID 列 '{id_column}'，将自动生成 ID")
        df_input[id_column] = [f"mol{i+1}" for i in range(len(df_input))]
        id_col_actual = id_column
    
    # 创建预测器配置
    config = PredictionConfig(
        batch_size=batch_size,
        n_workers=n_workers,
        device=device
    )
    
    # 初始化预测器
    print("初始化预测器...")
    predictor = ParallelBroadSpectrumPredictor(config)
    
    # 准备分子输入
    molecules = [
        MoleculeInput(smiles=row[smiles_col_actual], chem_id=str(row[id_col_actual]))
        for _, row in df_input.iterrows()
    ]
    
    # 执行预测
    print("开始预测...")
    results = predictor.predict_batch(molecules, include_strain_predictions=include_strain_predictions)
    
    # 转换结果为 DataFrame
    results_dicts = [r.to_dict() for r in results]
    df_results = pd.DataFrame(results_dicts)
    
    # 合并原始数据和预测结果
    # 使用 chem_id 作为键进行合并
    df_input['_merge_id'] = df_input[id_col_actual].astype(str)
    df_results['_merge_id'] = df_results['chem_id'].astype(str)
    
    df_output = df_input.merge(
        df_results.drop(columns=['chem_id']),
        on='_merge_id',
        how='left'
    )
    df_output = df_output.drop(columns=['_merge_id'])
    
    # 如果包含菌株级别预测，将其添加到输出中
    if include_strain_predictions:
        print("合并菌株级别预测数据...")
        # 收集所有菌株级别预测
        all_strain_predictions = []
        for result in results:
            if result.strain_predictions is not None and not result.strain_predictions.empty:
                all_strain_predictions.append(result.strain_predictions)
        
        if all_strain_predictions:
            # 合并所有菌株预测
            df_strain_predictions = pd.concat(all_strain_predictions, ignore_index=True)
            
            # 将聚合结果和菌株预测合并
            # 为了在同一个 CSV 中展示，我们使用重复行的方式
            # 每个分子的聚合结果会重复40次（每个菌株一次）
            df_output_expanded = df_output.merge(
                df_strain_predictions,
                left_on=df_output.columns[df_output.columns.get_loc(id_col_actual)],
                right_on='chem_id',
                how='left',
                suffixes=('', '_strain')
            )
            
            # 移除重复的 chem_id 列
            if 'chem_id_strain' in df_output_expanded.columns:
                df_output_expanded = df_output_expanded.drop(columns=['chem_id_strain'])
            
            df_output = df_output_expanded
    
    # 生成输出路径
    if output_path is None:
        if add_suffix:
            output_path = str(input_path_obj.parent / f"{input_path_obj.stem}_predicted{input_path_obj.suffix}")
        else:
            output_path = str(input_path_obj.parent / f"{input_path_obj.stem}_output{input_path_obj.suffix}")
    elif add_suffix:
        output_path_obj = Path(output_path)
        output_path = str(output_path_obj.parent / f"{output_path_obj.stem}_predicted{output_path_obj.suffix}")
    
    # 保存结果
    print(f"保存结果到: {output_path}")
    df_output.to_csv(output_path, index=False)
    
    print(f"完成! 预测了 {len(results)} 个分子")
    print(f"其中 {sum(r.broad_spectrum for r in results)} 个分子被预测为广谱抗菌")
    
    return df_output


def predict_multiple_files(
    input_paths: List[str],
    output_dir: Optional[str] = None,
    smiles_column: str = "smiles",
    id_column: str = "chem_id",
    batch_size: int = 100,
    n_workers: Optional[int] = None,
    device: str = "auto",
    add_suffix: bool = True,
    include_strain_predictions: bool = False
) -> List[pd.DataFrame]:
    """
    批量预测多个 CSV 文件
    
    Args:
        input_paths: 输入 CSV 文件路径列表
        output_dir: 输出目录，如果为 None 则在原文件目录生成
        smiles_column: SMILES 列名
        id_column: 化合物 ID 列名
        batch_size: 批处理大小
        n_workers: 工作进程数
        device: 计算设备
        add_suffix: 是否在输出文件名后添加预测后缀
        include_strain_predictions: 是否在输出中包含40种菌株的预测详情
    
    Returns:
        包含预测结果的 DataFrame 列表
    """
    
    results = []
    
    for input_path in input_paths:
        input_path_obj = Path(input_path)
        
        # 确定输出路径
        if output_dir is not None:
            output_dir_obj = Path(output_dir)
            output_dir_obj.mkdir(parents=True, exist_ok=True)
            
            if add_suffix:
                output_path = str(output_dir_obj / f"{input_path_obj.stem}_predicted{input_path_obj.suffix}")
            else:
                output_path = str(output_dir_obj / input_path_obj.name)
        else:
            output_path = None
        
        # 预测单个文件
        try:
            df_result = predict_csv_file(
                input_path=input_path,
                output_path=output_path,
                smiles_column=smiles_column,
                id_column=id_column,
                batch_size=batch_size,
                n_workers=n_workers,
                device=device,
                add_suffix=add_suffix,
                include_strain_predictions=include_strain_predictions
            )
            results.append(df_result)
        except Exception as e:
            print(f"处理文件 {input_path} 时出错: {e}")
            continue
    
    return results


# ============================================================================
# 命令行接口
# ============================================================================

@click.command()
@click.argument('input_path', type=click.Path(exists=True))
@click.argument('output_path', type=click.Path(), required=False)
@click.option('--smiles-column', '-s', default='smiles',
              help='SMILES 列名 (默认: smiles)')
@click.option('--id-column', '-i', default='chem_id',
              help='化合物 ID 列名 (默认: chem_id)')
@click.option('--batch-size', '-b', default=100, type=int,
              help='批处理大小 (默认: 100)')
@click.option('--n-workers', '-w', default=None, type=int,
              help='工作进程数 (默认: CPU 核心数)')
@click.option('--device', '-d', default='auto',
              type=click.Choice(['auto', 'cpu', 'cuda:0', 'cuda:1'], case_sensitive=False),
              help='计算设备 (默认: auto)')
@click.option('--add-suffix/--no-add-suffix', default=True,
              help='是否在输出文件名后添加 "_predicted" 后缀 (默认: 添加)')
@click.option('--include-strain-predictions', is_flag=True, default=False,
              help='在输出中包含40种菌株的详细预测数据（每个分子将产生40行数据，对应每个菌株的预测概率和抑制情况）')
def cli(input_path, output_path, smiles_column, id_column, batch_size, n_workers, device, add_suffix, include_strain_predictions):
    """
    使用 MolE 模型预测小分子 SMILES 的抗菌活性
    
    INPUT_PATH: 输入 CSV 文件路径
    
    OUTPUT_PATH: 输出 CSV 文件路径 (可选，默认在原文件目录生成)
    
    默认输出包含聚合的抗菌活性指标（广谱抗菌评分、抑制菌株数等）。
    使用 --include-strain-predictions 可以额外包含每个菌株的详细预测数据。
    
    示例:
    
        # 基本用法（仅输出聚合结果）
        python mole_predictor.py input.csv output.csv
        
        # 包含40种菌株的详细预测数据
        python mole_predictor.py input.csv output.csv --include-strain-predictions
        
        # 指定列名和设备
        python mole_predictor.py input.csv -s SMILES -i ID --device cuda:0
        
        # 自定义批处理大小
        python mole_predictor.py input.csv --device cuda:0 --batch-size 200
    """
    
    try:
        predict_csv_file(
            input_path=input_path,
            output_path=output_path,
            smiles_column=smiles_column,
            id_column=id_column,
            batch_size=batch_size,
            n_workers=n_workers,
            device=device,
            add_suffix=add_suffix,
            include_strain_predictions=include_strain_predictions
        )
    except Exception as e:
        click.echo(f"错误: {e}", err=True)
        sys.exit(1)


if __name__ == '__main__':
    cli()