# SIME - Structure-Informed Macrolide Expansion

SIME 是一个用于大环内酯类化合物结构扩展和抗菌活性预测的工具。

## 目录

- [原有功能](#原有功能)
- [MolE 抗菌活性预测](#mole-抗菌活性预测)
  - [快速开始](#快速开始)
  - [安装依赖](#安装依赖)
  - [使用方法](#使用方法)
  - [输出说明](#输出说明)
- [项目结构](#项目结构)
- [常见问题](#常见问题)

---

## 原有功能

SIME 提供大环内酯类化合物的结构设计和合成路径分析功能。

---

## MolE 抗菌活性预测

本工具集成了 MolE（Molecular Embeddings）模型，可以预测小分子的广谱抗菌活性。

### 快速开始

#### 使用 uv（推荐）

```bash
# 1. 创建虚拟环境（Python 3.12）
uv venv --python 3.12 --seed .venv

# 2. 激活环境
source .venv/bin/activate  # Linux/Mac
# 或
.venv\Scripts\activate  # Windows

# 3. 使用 uv 安装依赖
uv pip install -r requirements-mole.txt

# 4. 验证安装
python verify_setup.py

# 5. 运行预测
python utils/mole_predictor.py Data/fragment/Frags-Enamine-18M.csv
```

#### 使用 pyproject.toml 配置（uv 推荐）

项目提供了两个环境配置：

1. **SIME 原始环境** - 用于大环内酯结构设计

```bash
# 使用 uv 创建默认环境
uv sync
```

2. **MolE 预测环境** - 用于抗菌活性预测

```bash
# 使用 uv 创建 MolE 环境
uv sync --extra mole
```

#### 使用 pixi 配置（conda 用户推荐）

如果你使用 conda 或需要更好的包管理，可以使用 pixi：

```bash
# 安装 pixi（如果还没有）
curl -fsSL https://pixi.sh/install.sh | bash

# 创建 SIME 原始环境
pixi install

# 创建 MolE 预测环境
pixi install -e mole

# 激活 MolE 环境
pixi shell -e mole

# 在 pixi 环境中运行预测
pixi run -e mole predict Data/fragment/test_100.csv
```

### 安装依赖

#### 方法 1: 使用 uv（推荐）

```bash
# 创建虚拟环境
uv venv --python 3.12 .venv
source .venv/bin/activate

# 安装依赖
uv pip install -r requirements-mole.txt
```

#### 方法 2: 使用 pixi

```bash
# 创建虚拟环境
pixi init

# 基础环境
pixi add python=3.12

# nvidia cuda工具链
pixi workspace channel add nvidia
pixi add nvidia::cuda-toolkit=12.8

# 科学计算 安装 pandas 会自动安装上 numpy
pixi add 

# torch-geometric
pixi add conda-forge::pandas conda-forge::torch-geometric conda-forge::xgboost conda-forge::pyyaml conda-forge::rdkit conda-forge::pip conda-forge::click conda-forge::openpyxl

# PyTorch相关（指定通道）
# 1. 添加 pytorch 频道  conda 太旧改为使用 pypi
# pixi workspace channel add pytorch
# pixi add pytorch::pytorch=2.6 pytorch::pytorch-cuda=12.4
pixi add --pypi torch==2.8.0 torchvision==0.23.0 torchaudio==2.8.0

# 然后在 pixi.toml 中手动编辑为:
[pypi-dependencies]  
torch = { version = "==2.8.0", index = "https://download.pytorch.org/whl/cu128" }  
torchvision = { version = "==0.23.0", index = "https://download.pytorch.org/whl/cu128" }  
torchaudio = { version = "==2.8.0", index = "https://download.pytorch.org/whl/cu128" }

# 安装依赖
pixi install

# 激活
pixi shell
```

不同机器使用配置方式：

```bash
# 在 Linux GPU 机器上安装和运行  
pixi install --environment gpu  
pixi run --environment gpu <your-task>  
  
# 在 macOS 或 CPU 机器上安装和运行  
pixi install --environment cpu  
pixi run --environment cpu <your-task>  
  
# 或使用默认环境(CPU)  
pixi install  
pixi run <your-task>
```

#### RDKit 安装建议

RDKit 推荐使用 conda 安装：

```bash
conda install -c conda-forge rdkit
```

### 使用方法

#### 1. 命令行使用

**基本用法：**

```bash
# 预测 CSV 文件（仅聚合结果）
python utils/mole_predictor.py input.csv

# 包含 40 种菌株的详细预测数据
python utils/mole_predictor.py input.csv --include-strain-predictions

# 指定输出路径
python utils/mole_predictor.py input.csv output.csv

# 自定义列名
python utils/mole_predictor.py input.csv output.csv \
  --smiles-column SMILES \
  --id-column compound_id

# 使用 GPU 加速
python utils/mole_predictor.py input.csv --device cuda:0

# 调整批次大小和工作进程
python utils/mole_predictor.py input.csv \
  --batch-size 200 \
  --n-workers 8

# 完整示例：包含菌株预测 + GPU 加速
python utils/mole_predictor.py input.csv output.csv \
  --include-strain-predictions \
  --device cuda:0 \
  --batch-size 200
```

**查看所有选项：**

```bash
python utils/mole_predictor.py --help
```

**预测项目数据：**

```bash
# 预测 Frags-Enamine-18M.csv
# 创建测试文件（前 1001 行，包含表头）
head -1001 Data/fragment/Frags-Enamine-18M.csv > Data/fragment/test_1000.csv

# 测试命令 - 保守参数
nohup pixi run python utils/mole_predictor.py \
  Data/fragment/test_1000.csv \
  Data/fragment/test_1000_predicted.csv \
  --device cuda:0 \
  --batch-size 100 \
  --n-workers 8 \
  > Data/fragment/mole_test_1000.log 2>&1 &

# 查看日志
tail -f Data/fragment/mole_test_1000.log

# 预测 GDB11-27M.csv
python utils/mole_predictor.py Data/fragment/GDB11-27M.csv
```

#### 2. Python API 使用

**预测单个文件：**

```python
from utils.mole_predictor import predict_csv_file

# 基本使用（仅聚合结果）
df_result = predict_csv_file(
    input_path="Data/fragment/Frags-Enamine-18M.csv",
    output_path="results/predictions.csv",
    smiles_column="smiles",
    batch_size=100,
    device="auto"
)

# 包含 40 种菌株的详细预测数据
df_result_with_strains = predict_csv_file(
    input_path="Data/fragment/Frags-Enamine-18M.csv",
    output_path="results/predictions_with_strains.csv",
    smiles_column="smiles",
    batch_size=100,
    device="auto",
    include_strain_predictions=True  # 启用菌株级别预测
)

# 查看结果
print(f"总分子数: {len(df_result)}")
print(f"广谱分子数: {df_result['broad_spectrum'].sum()}")

# 如果包含菌株预测，数据行数会增加（每个分子 40 行）
if 'strain_name' in df_result_with_strains.columns:
    print(f"包含菌株预测的总行数: {len(df_result_with_strains)}")
    print(f"菌株数: {df_result_with_strains['strain_name'].nunique()}")
```

**批量预测多个文件：**

```python
from utils.mole_predictor import predict_multiple_files

input_files = [
    "Data/fragment/Frags-Enamine-18M.csv",
    "Data/fragment/GDB11-27M.csv"
]

results = predict_multiple_files(
    input_paths=input_files,
    output_dir="results/",
    smiles_column="smiles",
    batch_size=100,
    device="auto"
)
```

**直接使用预测器：**

```python
from models import (
    ParallelBroadSpectrumPredictor,
    PredictionConfig,
    MoleculeInput
)

# 创建配置
config = PredictionConfig(
    batch_size=100,
    device="auto"  # 或 "cpu", "cuda:0"
)

# 创建预测器
predictor = ParallelBroadSpectrumPredictor(config)

# 预测单个分子（仅聚合结果）
molecule = MoleculeInput(smiles="CCO", chem_id="ethanol")
result = predictor.predict_single(molecule)

print(f"化合物ID: {result.chem_id}")
print(f"广谱抗菌: {result.broad_spectrum}")
print(f"抗菌得分: {result.apscore_total:.3f}")
print(f"抑制菌株数: {result.ginhib_total}")

# 批量预测（包含 40 种菌株的详细预测）
smiles_list = ["CCO", "c1ccccc1", "CC(=O)O"]
chem_ids = ["ethanol", "benzene", "acetic_acid"]

results = predictor.predict_batch(
    [MoleculeInput(smiles=s, chem_id=c) for s, c in zip(smiles_list, chem_ids)],
    include_strain_predictions=True  # 启用菌株级别预测
)

for r in results:
    print(f"\n{r.chem_id}:")
    print(f"  广谱抗菌: {r.broad_spectrum}")
    print(f"  抗菌得分: {r.apscore_total:.3f}")
    print(f"  抑制菌株数: {r.ginhib_total}")
    
    # 访问菌株级别预测数据（DataFrame 格式）
    if r.strain_predictions is not None:
        print(f"  菌株预测数据形状: {r.strain_predictions.shape}")
        print(f"  示例菌株预测:")
        print(r.strain_predictions.head(3))
        
        # 强化学习场景：提取特定菌株的预测概率
        strain_probs = r.strain_predictions['antimicrobial_predictive_probability'].values
        print(f"  预测概率向量形状: {strain_probs.shape}")  # (40,)
        
        # 或转换为类型安全的列表
        strain_list = r.to_strain_predictions_list()
        print(f"  第一个菌株: {strain_list[0].strain_name}")
        print(f"  预测概率: {strain_list[0].antimicrobial_predictive_probability:.6f}")
```

### 输出说明

#### 1. 聚合结果（默认输出）

预测结果会添加以下 7 个新列：

| 列名 | 类型 | 说明 |
|------|------|------|
| `apscore_total` | float | 总体抗菌潜力分数（对数尺度，值越大抗菌活性越强） |
| `apscore_gnegative` | float | 革兰阴性菌抗菌潜力分数 |
| `apscore_gpositive` | float | 革兰阳性菌抗菌潜力分数 |
| `ginhib_total` | int | 被抑制的菌株总数（0-40） |
| `ginhib_gnegative` | int | 被抑制的革兰阴性菌株数 |
| `ginhib_gpositive` | int | 被抑制的革兰阳性菌株数 |
| `broad_spectrum` | int | 是否为广谱抗菌（1=是，0=否） |

**输出示例**（每个分子一行）：

```csv
smiles,chem_id,apscore_total,apscore_gnegative,apscore_gpositive,ginhib_total,ginhib_gnegative,ginhib_gpositive,broad_spectrum
CCO,mol1,-9.93,-10.17,-9.74,0,0,0,0
```

#### 2. 菌株级别预测（使用 `--include-strain-predictions` 时）

启用菌株级别预测后，输出会包含以下额外列（每个分子 40 行）：

| 列名 | 类型 | 说明 |
|------|------|------|
| `pred_id` | str | 预测ID，格式为 `chem_id:strain_name` |
| `strain_name` | str | 菌株名称（40 种菌株之一） |
| `antimicrobial_predictive_probability` | float | XGBoost 预测的抗菌概率（0-1） |
| `no_growth_probability` | float | 预测不抑制的概率（1 - antimicrobial_predictive_probability） |
| `growth_inhibition` | int | 二值化抑制结果（0=不抑制，1=抑制） |
| `gram_stain` | str | 革兰染色类型（negative 或 positive） |

**输出示例**（每个分子 40 行，对应 40 个菌株）：

```csv
smiles,chem_id,apscore_total,...,pred_id,strain_name,antimicrobial_predictive_probability,no_growth_probability,growth_inhibition,gram_stain
CCO,mol1,-9.93,...,mol1:Akkermansia muciniphila (NT5021),Akkermansia muciniphila (NT5021),0.000102,0.999898,0,negative
CCO,mol1,-9.93,...,mol1:Bacteroides caccae (NT5050),Bacteroides caccae (NT5050),0.000155,0.999845,0,negative
...（共 40 行，对应 40 个菌株）
```

**数据使用场景**：

1. **强化学习状态表示**：
   ```python
   # 提取预测概率作为状态向量
   state_vector = result.strain_predictions['antimicrobial_predictive_probability'].values
   # 形状: (40,) - 可直接用于 RL 环境
   ```

2. **筛选特定菌株**：
   ```python
   # 筛选革兰阴性菌
   gram_negative = result.strain_predictions[
       result.strain_predictions['gram_stain'] == 'negative'
   ]
   ```

3. **可视化分析**：
   ```python
   import matplotlib.pyplot as plt
   df = result.strain_predictions
   df.plot(x='strain_name', y='antimicrobial_predictive_probability', kind='bar')
   ```

#### 40 种测试菌株列表

预测涵盖以下 40 种人类肠道菌株：

**革兰阴性菌（23 种）**：
- Akkermansia muciniphila (NT5021)
- Bacteroides 属: caccae, fragilis (ET/NT), ovatus, thetaiotaomicron, uniformis, vulgatus, xylanisolvens (8 种)
- Escherichia coli 各亚型 (4 种)
- Klebsiella pneumoniae (NT5049)
- 其他肠道革兰阴性菌

**革兰阳性菌（17 种）**：
- Bifidobacterium 属 (3 种)
- Clostridium 属 (5 种)
- Enterococcus 属 (2 种)
- Lactobacillus 属 (3 种)
- Streptococcus 属 (2 种)
- 其他肠道革兰阳性菌

完整菌株列表详见 `Data/mole/README.md`。

#### 广谱抗菌判断标准

默认情况下，如果一个分子能抑制 **10 个或更多菌株** (`ginhib_total >= 10`)，则被认为是广谱抗菌分子。

#### 输出文件位置

默认情况下，输出文件会添加 `_predicted` 后缀：

- 输入: `Data/fragment/Frags-Enamine-18M.csv`
- 输出: `Data/fragment/Frags-Enamine-18M_predicted.csv`
- 输出（含菌株）: `Data/fragment/Frags-Enamine-18M_predicted.csv`（每个分子 40 行）

#### 数据量说明

- **仅聚合结果**：输出行数 = 输入分子数
- **包含菌株预测**：输出行数 = 输入分子数 × 40

---

## 项目结构

```
SIME/
├── models/                          # MolE 预测模型
│   ├── __init__.py
│   ├── broad_spectrum_predictor.py  # 核心预测器
│   ├── dataset_representation.py    # 数据集表示
│   ├── ginet_concat.py             # GIN 神经网络
│   └── mole_representation.py      # MolE 表示生成
│
├── utils/
│   ├── mole_predictor.py           # 预测工具脚本
│   └── ... (其他工具)
│
├── Data/
│   └── fragment/                    # 待预测数据
│       ├── Frags-Enamine-18M.csv
│       └── GDB11-27M.csv
│
├── pyproject.toml                   # uv 项目配置
├── requirements.txt                 # SIME 原始依赖
├── requirements-mole.txt           # MolE 预测依赖
│
├── verify_setup.py                 # 设置验证工具
├── check_mole_dependencies.py      # 依赖检查工具
└── test_mole_predictor.py          # 功能测试
```

---

## 依赖说明

### SIME 原始依赖 (requirements.txt)

用于大环内酯结构设计功能。

### MolE 预测依赖 (requirements-mole.txt)

用于抗菌活性预测，主要包括：

- **深度学习**: torch, torch-geometric
- **科学计算**: numpy, pandas, scipy
- **机器学习**: scikit-learn, xgboost
- **化学信息**: rdkit
- **其他**: openpyxl, pyyaml, click

---

## 验证和测试

### 验证安装

```bash
# 检查 Python 依赖
python verify_setup.py

# 检查模型文件
python check_mole_dependencies.py
```

### 运行测试

```bash
# 功能测试（使用小规模测试数据）
python test_mole_predictor.py
```

---

## 常见问题

### Q1: 如何处理大文件？

**方案 1：** 增加批次大小和工作进程数

```bash
python utils/mole_predictor.py large_file.csv \
  --batch-size 500 \
  --n-workers 8
```

**方案 2：** 先提取部分数据测试

```bash
# 提取前 1000 行
head -1001 large_file.csv > test_1000.csv
python utils/mole_predictor.py test_1000.csv
```

### Q2: 如何只使用 CPU？

```bash
python utils/mole_predictor.py input.csv --device cpu
```

### Q3: 列名大小写问题？

工具会自动进行大小写不敏感的列名匹配，所以 `SMILES`、`smiles`、`Smiles` 都可以识别。

### Q4: ModuleNotFoundError 错误？

确保已安装依赖：

```bash
uv pip install -r requirements-mole.txt
```

对于 RDKit，推荐使用 conda：

```bash
conda install -c conda-forge rdkit
```

### Q5: 如何自定义模型路径？

```python
from models import PredictionConfig, ParallelBroadSpectrumPredictor

config = PredictionConfig(
    xgboost_model_path="/path/to/model.pkl",
    mole_model_path="/path/to/mole_model",
    strain_categories_path="/path/to/strain_data.tsv.gz",
    gram_info_path="/path/to/gram_info.xlsx",
    app_threshold=0.044,
    min_nkill=10,
    batch_size=100,
    device="auto"
)

predictor = ParallelBroadSpectrumPredictor(config)
```

### Q6: GPU 内存不足？

减小批次大小：

```bash
python utils/mole_predictor.py input.csv --batch-size 50
```

### Q7: 模型文件在哪里？

模型文件位于相邻的 `mole_broad_spectrum_parallel` 项目中：

```
../mole_broad_spectrum_parallel/
├── pretrained_model/model_ginconcat_btwin_100k_d8000_l0.0001/
│   ├── config.yaml
│   └── model.pth
├── data/03.model_evaluation/MolE-XGBoost-08.03.2024_14.20.pkl
└── ...
```

运行 `python check_mole_dependencies.py` 检查文件是否存在。

---

## 性能建议

- **使用 GPU**: 设置 `--device cuda:0` 可大幅加速（需要 CUDA）
- **调整批次**: 较大的批次（100-500）通常更快
- **多进程**: 使用 `--n-workers` 指定工作进程数
- **首次加载**: 首次运行需要加载模型（~30秒），后续会更快

### 性能参考

| 分子数量 | CPU (8核) | GPU (CUDA) |
|---------|----------|------------|
| 100     | ~30秒    | ~10秒      |
| 1,000   | ~5分钟   | ~1分钟     |
| 10,000  | ~50分钟  | ~8分钟     |

---

## 系统要求

- **Python**: 3.7 或更高版本（推荐 3.12）
- **内存**: 最低 8 GB RAM
- **存储**: 至少 2 GB 可用空间
- **GPU**: 可选，但强烈推荐（需要 CUDA 支持）

---

## 技术支持

如有问题：

1. 查看验证结果: `python verify_setup.py`
2. 检查模型文件: `python check_mole_dependencies.py`
3. 运行功能测试: `python test_mole_predictor.py`

---

## 许可

详见 LICENSE 文件。

## 引用

如果使用本工具，请引用相关论文。

---

**更新日期**: 2025-10-16  
**版本**: 1.0.0