327 lines
11 KiB
Python
Executable File
327 lines
11 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
"""
|
|
Functional Annotation Skill (API Mode) - 功能注释分析技能（API模式）
|
|
使用 eggNOG API 在线获取 GO 注释，无需下载数据库
|
|
|
|
Usage:
|
|
python run_api.py --input protein.fa --output ./annotation
|
|
python run_api.py --demo --output ./annotation
|
|
|
|
特点：
|
|
- 无需下载数据库
|
|
- 适合小数据集（<1000基因）
|
|
- 速度较慢（约1秒/基因）
|
|
"""
|
|
|
|
import argparse
|
|
import hashlib
|
|
import json
|
|
import logging
|
|
import os
|
|
import sys
|
|
import time
|
|
from collections import defaultdict
|
|
from datetime import datetime
|
|
from pathlib import Path
|
|
from typing import Dict, List, Optional, Set
|
|
|
|
import requests
|
|
import yaml
|
|
|
|
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# eggNOG API endpoints
|
|
EGGNOG_SCAN_URL = "http://eggnogapi5.embl.de/fast_webscan"
|
|
EGGNOG_GO_URL = "http://eggnogapi5.embl.de/nog_data/json/go_terms"
|
|
|
|
HEADERS = {
|
|
"User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36",
|
|
"Accept": "application/json"
|
|
}
|
|
|
|
|
|
class AnnotationAPISkill:
|
|
"""功能注释API技能主类"""
|
|
|
|
def __init__(self, config_path: Optional[str] = None):
|
|
self.config = self._load_config(config_path)
|
|
self.commands = []
|
|
self.checksums = {}
|
|
self.stats = {}
|
|
|
|
def _load_config(self, config_path: Optional[str]) -> Dict:
|
|
if config_path and os.path.exists(config_path):
|
|
with open(config_path) as f:
|
|
return yaml.safe_load(f)
|
|
return {
|
|
"parameters": {
|
|
"api": {
|
|
"delay": 1.0,
|
|
"timeout": 60,
|
|
"db": "euk"
|
|
}
|
|
}
|
|
}
|
|
|
|
def _compute_checksum(self, filepath: Path) -> str:
|
|
sha256 = hashlib.sha256()
|
|
with open(filepath, "rb") as f:
|
|
for chunk in iter(lambda: f.read(8192), b""):
|
|
sha256.update(chunk)
|
|
return sha256.hexdigest()
|
|
|
|
def _record_file(self, filepath: Path, file_type: str = "output"):
|
|
if filepath.exists():
|
|
self.checksums[str(filepath)] = {
|
|
"sha256": self._compute_checksum(filepath),
|
|
"size": filepath.stat().st_size,
|
|
"type": file_type
|
|
}
|
|
|
|
def read_fasta(self, fasta_file: Path) -> List[tuple]:
|
|
"""读取FASTA文件"""
|
|
sequences = []
|
|
with open(fasta_file) as f:
|
|
name, seq = None, []
|
|
for line in f:
|
|
line = line.strip()
|
|
if line.startswith('>'):
|
|
if name:
|
|
sequences.append((name, ''.join(seq)))
|
|
name = line[1:].split()[0]
|
|
seq = []
|
|
else:
|
|
seq.append(line)
|
|
if name:
|
|
sequences.append((name, ''.join(seq)))
|
|
return sequences
|
|
|
|
def submit_to_eggnog(self, session: requests.Session, fasta_seq: str, db: str = "euk") -> Optional[str]:
|
|
"""提交序列到eggNOG，返回nogname"""
|
|
payload = {"fasta": fasta_seq, "nqueries": 1, "db": db}
|
|
try:
|
|
resp = session.post(EGGNOG_SCAN_URL, json=payload, headers=HEADERS, timeout=60)
|
|
if resp.status_code == 200:
|
|
data = resp.json()
|
|
if 'seq_matches' in data and data['seq_matches']:
|
|
result = data['seq_matches'][0]
|
|
if 'hit' in result and 'matches' in result['hit']:
|
|
return result['hit']['matches'][0].get('nogname')
|
|
except Exception as e:
|
|
logger.warning(f"提交错误: {e}")
|
|
return None
|
|
|
|
def get_go_terms(self, session: requests.Session, nogname: str) -> Set[str]:
|
|
"""获取GO terms"""
|
|
if not nogname:
|
|
return set()
|
|
url = f"{EGGNOG_GO_URL}/{nogname}"
|
|
try:
|
|
resp = session.get(url, headers=HEADERS, timeout=30)
|
|
if resp.status_code == 200:
|
|
data = resp.json()
|
|
go_ids = set()
|
|
for category, terms in data.get('go_terms', {}).items():
|
|
for term in terms:
|
|
go_id = term[0]
|
|
if go_id.startswith('GO:'):
|
|
go_ids.add(go_id)
|
|
return go_ids
|
|
except Exception as e:
|
|
logger.warning(f"GO获取错误: {e}")
|
|
return set()
|
|
|
|
def run_api_annotation(self, protein_fasta: Path, output_dir: Path,
|
|
prefix: str = "annotation") -> Dict[str, Path]:
|
|
"""运行API注释"""
|
|
params = self.config["parameters"]["api"]
|
|
delay = params.get("delay", 1.0)
|
|
db = params.get("db", "euk")
|
|
|
|
output_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
# 读取序列
|
|
sequences = self.read_fasta(protein_fasta)
|
|
total = len(sequences)
|
|
logger.info(f"[API] 共 {total} 条序列")
|
|
|
|
if total > 1000:
|
|
logger.warning(f"序列数较多({total})，API模式可能较慢，建议使用本地emapper")
|
|
|
|
session = requests.Session()
|
|
gene2go = defaultdict(set)
|
|
failed_genes = []
|
|
|
|
for i, (gene_id, seq) in enumerate(sequences):
|
|
logger.info(f"[{i+1}/{total}] {gene_id}")
|
|
|
|
fasta = f">{gene_id}\n{seq}"
|
|
|
|
# 1. 获取nogname
|
|
nogname = self.submit_to_eggnog(session, fasta, db)
|
|
if not nogname:
|
|
logger.warning(f" 获取nogname失败: {gene_id}")
|
|
failed_genes.append(gene_id)
|
|
time.sleep(delay)
|
|
continue
|
|
|
|
# 2. 获取GO terms
|
|
go_ids = self.get_go_terms(session, nogname)
|
|
if go_ids:
|
|
gene2go[gene_id] = go_ids
|
|
logger.info(f" 找到 {len(go_ids)} 个GO terms")
|
|
|
|
time.sleep(delay)
|
|
|
|
session.close()
|
|
|
|
# 保存结果
|
|
annot_file = output_dir / f"{prefix}.gene2go.txt"
|
|
with open(annot_file, 'w') as f:
|
|
for gene_id in sorted(gene2go.keys()):
|
|
for go_id in sorted(gene2go[gene_id]):
|
|
f.write(f"{gene_id}\t{go_id}\n")
|
|
|
|
self._record_file(annot_file)
|
|
|
|
# 统计
|
|
all_go_terms = set()
|
|
for gos in gene2go.values():
|
|
all_go_terms.update(gos)
|
|
|
|
self.stats = {
|
|
"total_genes": total,
|
|
"annotated_genes": len(gene2go),
|
|
"failed_genes": len(failed_genes),
|
|
"go_term_count": len(all_go_terms),
|
|
}
|
|
|
|
return {
|
|
"gene2go": annot_file,
|
|
}
|
|
|
|
def generate_report(self, output_dir: Path) -> Path:
|
|
"""生成报告"""
|
|
report_file = output_dir / "report.md"
|
|
|
|
report_content = f"""# 功能注释分析报告 (API模式)
|
|
|
|
## 基本信息
|
|
|
|
- **分析方式**: eggNOG API (在线)
|
|
- **分析时间**: {datetime.now().strftime("%Y-%m-%d %H:%M:%S")}
|
|
- **API延迟**: {self.config["parameters"]["api"].get("delay", 1.0)}s
|
|
|
|
## 统计结果
|
|
|
|
| 指标 | 数量 |
|
|
|------|------|
|
|
| 总基因数 | {self.stats.get('total_genes', 0)} |
|
|
| 成功注释 | {self.stats.get('annotated_genes', 0)} |
|
|
| 失败基因 | {self.stats.get('failed_genes', 0)} |
|
|
| 独特GO terms | {self.stats.get('go_term_count', 0)} |
|
|
|
|
## 输出文件
|
|
|
|
| 文件 | 描述 |
|
|
|------|------|
|
|
| *.gene2go.txt | 基因-GO对应表 |
|
|
|
|
## 注意事项
|
|
|
|
API模式适合小数据集，大数据集建议使用本地eggNOG-mapper。
|
|
"""
|
|
|
|
with open(report_file, "w") as f:
|
|
f.write(report_content)
|
|
|
|
self._record_file(report_file)
|
|
return report_file
|
|
|
|
def run(self, protein_fasta: Path, output_dir: Path, prefix: str = "annotation") -> Dict:
|
|
"""运行完整流程"""
|
|
logger.info(f"Starting API annotation for {protein_fasta}")
|
|
|
|
self._record_file(protein_fasta, "input")
|
|
|
|
results = self.run_api_annotation(protein_fasta, output_dir, prefix)
|
|
|
|
report = self.generate_report(output_dir)
|
|
|
|
logger.info(f"API annotation completed!")
|
|
logger.info(f" - Annotated: {self.stats['annotated_genes']}/{self.stats['total_genes']}")
|
|
logger.info(f" - GO terms: {self.stats['go_term_count']}")
|
|
|
|
return {
|
|
"status": "success",
|
|
"output_dir": str(output_dir),
|
|
"report": str(report),
|
|
"gene2go": str(results["gene2go"]),
|
|
"stats": self.stats,
|
|
}
|
|
|
|
def run_demo(self, output_dir: Path) -> Dict:
|
|
"""Demo模式"""
|
|
logger.info("Running API demo mode")
|
|
|
|
demo_dir = output_dir / "demo_data"
|
|
demo_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
# 创建模拟蛋白文件
|
|
protein_fasta = demo_dir / "demo_proteins.fa"
|
|
protein_fasta.write_text(">Gene001\nMKTLLLPLLLLVLLLLLVQGSISAQAAAAAA\n"
|
|
">Gene002\nMKTLLLPLLLLVLLLLLVQGSISAQBBBBBB\n")
|
|
|
|
output_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
# 模拟结果
|
|
gene2go_file = output_dir / "annotation.gene2go.txt"
|
|
gene2go_file.write_text("Gene001\tGO:0005524\nGene001\tGO:0004674\nGene002\tGO:0003677\n")
|
|
|
|
self.stats = {
|
|
"total_genes": 2,
|
|
"annotated_genes": 2,
|
|
"failed_genes": 0,
|
|
"go_term_count": 3,
|
|
}
|
|
|
|
report = self.generate_report(output_dir)
|
|
|
|
return {
|
|
"status": "success",
|
|
"output_dir": str(output_dir),
|
|
"report": str(report),
|
|
"gene2go": str(gene2go_file),
|
|
"stats": self.stats,
|
|
}
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description="Functional Annotation Skill (API Mode)")
|
|
parser.add_argument("--input", "-i", type=Path, help="蛋白序列FASTA文件")
|
|
parser.add_argument("--output", "-o", type=Path, required=True, help="输出目录")
|
|
parser.add_argument("--prefix", type=str, default="annotation", help="输出文件前缀")
|
|
parser.add_argument("--delay", type=float, default=1.0, help="API请求间隔(秒)")
|
|
parser.add_argument("--demo", action="store_true", help="运行Demo模式")
|
|
|
|
args = parser.parse_args()
|
|
|
|
skill = AnnotationAPISkill()
|
|
skill.config["parameters"]["api"]["delay"] = args.delay
|
|
|
|
if args.demo:
|
|
result = skill.run_demo(args.output)
|
|
else:
|
|
if not args.input:
|
|
parser.error("--input is required (unless using --demo)")
|
|
result = skill.run(args.input, args.output, args.prefix)
|
|
|
|
print(json.dumps(result, indent=2, ensure_ascii=False))
|
|
|
|
return 0 if result["status"] == "success" else 1
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(main())
|