391 lines
15 KiB
Python
Executable File
391 lines
15 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
"""
|
|
Functional Annotation Skill - 功能注释分析技能
|
|
使用 eggNOG-mapper 对蛋白序列进行功能注释，生成 .emapper.annotations 文件
|
|
|
|
Usage:
|
|
python run.py --input protein.fa --output ./annotation
|
|
python run.py --config config.yaml --output ./annotation
|
|
python run.py --demo --output ./annotation
|
|
"""
|
|
|
|
import argparse
|
|
import hashlib
|
|
import json
|
|
import logging
|
|
import os
|
|
import shutil
|
|
import subprocess
|
|
import sys
|
|
from datetime import datetime
|
|
from pathlib import Path
|
|
from typing import Dict, List, Optional, Tuple
|
|
|
|
import yaml
|
|
|
|
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class AnnotationSkill:
|
|
"""功能注释分析技能主类"""
|
|
|
|
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)
|
|
default_config = Path(__file__).parent / "config.yaml"
|
|
with open(default_config) as f:
|
|
return yaml.safe_load(f)
|
|
|
|
def _which(self, cmd: str) -> Optional[str]:
|
|
return shutil.which(cmd)
|
|
|
|
def _find_emapper(self) -> Tuple[str, str]:
|
|
"""
|
|
查找 emapper.py 路径和对应的 python 解释器
|
|
返回: (python_path, emapper_path)
|
|
"""
|
|
emapper = self._which("emapper.py")
|
|
if emapper:
|
|
# 使用系统默认的 python
|
|
return sys.executable, emapper
|
|
|
|
# 检查常见conda环境路径
|
|
conda_envs = [
|
|
Path.home() / "miniconda3/envs/test4py311",
|
|
Path.home() / "anaconda3/envs/eggnog",
|
|
Path.home() / "miniconda3/envs/eggnog",
|
|
]
|
|
for env_path in conda_envs:
|
|
emapper_path = env_path / "bin/emapper.py"
|
|
python_path = env_path / "bin/python"
|
|
if emapper_path.exists() and python_path.exists():
|
|
return str(python_path), str(emapper_path)
|
|
|
|
raise RuntimeError("emapper.py not found. Install: conda install -c bioconda eggnog-mapper")
|
|
|
|
def _run_command(self, cmd: List[str], description: str, check: bool = True):
|
|
logger.info(f"Running: {description}")
|
|
logger.info(f"Command: {' '.join(cmd)}")
|
|
|
|
self.commands.append({
|
|
"description": description,
|
|
"command": " ".join(cmd),
|
|
"timestamp": datetime.now().isoformat()
|
|
})
|
|
|
|
result = subprocess.run(cmd, capture_output=True, text=True, check=False)
|
|
|
|
if result.returncode != 0 and check:
|
|
logger.error(f"Command failed: {result.stderr}")
|
|
raise subprocess.CalledProcessError(result.returncode, cmd, output=result.stdout, stderr=result.stderr)
|
|
|
|
return result
|
|
|
|
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 _parse_annotations(self, annot_file: Path) -> Dict:
|
|
stats = {
|
|
"total_genes": 0, "with_go": 0, "with_kegg": 0,
|
|
"with_cog": 0, "with_pfam": 0,
|
|
"go_terms": set(), "kegg_terms": set(), "cog_categories": set(),
|
|
}
|
|
|
|
if not annot_file.exists():
|
|
return stats
|
|
|
|
with open(annot_file) as f:
|
|
header = None
|
|
for line in f:
|
|
if line.startswith("##"):
|
|
continue
|
|
if line.startswith("#"):
|
|
header = line.strip().lstrip("#").split("\t")
|
|
continue
|
|
|
|
if header:
|
|
parts = line.strip().split("\t")
|
|
if len(parts) < 10:
|
|
continue
|
|
|
|
stats["total_genes"] += 1
|
|
|
|
try:
|
|
go_idx = header.index("GOs") if "GOs" in header else -1
|
|
kegg_idx = header.index("KEGG_ko") if "KEGG_ko" in header else -1
|
|
cog_idx = header.index("COG_category") if "COG_category" in header else -1
|
|
pfam_idx = header.index("PFAMs") if "PFAMs" in header else -1
|
|
except ValueError:
|
|
continue
|
|
|
|
if go_idx >= 0 and go_idx < len(parts) and parts[go_idx] and parts[go_idx] != "-":
|
|
stats["with_go"] += 1
|
|
stats["go_terms"].update(parts[go_idx].split(","))
|
|
|
|
if kegg_idx >= 0 and kegg_idx < len(parts) and parts[kegg_idx] and parts[kegg_idx] != "-":
|
|
stats["with_kegg"] += 1
|
|
stats["kegg_terms"].update(parts[kegg_idx].split(","))
|
|
|
|
if cog_idx >= 0 and cog_idx < len(parts) and parts[cog_idx] and parts[cog_idx] != "-":
|
|
stats["with_cog"] += 1
|
|
stats["cog_categories"].update(parts[cog_idx])
|
|
|
|
if pfam_idx >= 0 and pfam_idx < len(parts) and parts[pfam_idx] and parts[pfam_idx] != "-":
|
|
stats["with_pfam"] += 1
|
|
|
|
stats["go_count"] = len(stats.pop("go_terms"))
|
|
stats["kegg_count"] = len(stats.pop("kegg_terms"))
|
|
stats["cog_count"] = len(stats.pop("cog_categories"))
|
|
|
|
return stats
|
|
|
|
def run_emapper(self, protein_fasta: Path, output_dir: Path, prefix: str = "annotation") -> Dict[str, Path]:
|
|
params = self.config["parameters"]["emapper"]
|
|
cpu = params.get("cpu", 8)
|
|
|
|
python_path, emapper_path = self._find_emapper()
|
|
logger.info(f"Using python: {python_path}")
|
|
logger.info(f"Using emapper: {emapper_path}")
|
|
|
|
output_dir.mkdir(parents=True, exist_ok=True)
|
|
output_prefix = output_dir / prefix
|
|
|
|
cmd = [
|
|
python_path, emapper_path,
|
|
"-i", str(protein_fasta),
|
|
"--output", str(output_prefix),
|
|
"-m", params.get("mode", "diamond"),
|
|
"--cpu", str(cpu),
|
|
"--evalue", str(params.get("evalue", 0.001)),
|
|
"--score", str(params.get("score", 60)),
|
|
"--query-cover", str(params.get("query_cover", 20)),
|
|
"--subject-cover", str(params.get("subject_cover", 20)),
|
|
"--tax-scope", params.get("tax_scope", "auto"),
|
|
"--go-evidence", params.get("go_evidence", "non-electronic"),
|
|
"--target-orthologs", params.get("target_orthologs", "all"),
|
|
]
|
|
|
|
# 添加数据库目录参数（如果指定）
|
|
if params.get("data_dir"):
|
|
cmd.extend(["--data_dir", str(params["data_dir"])])
|
|
|
|
if params.get("block_size"):
|
|
cmd.extend(["--block-size", str(params["block_size"])])
|
|
|
|
self._run_command(cmd, f"eggNOG-mapper annotation ({prefix})")
|
|
|
|
annot_file = output_dir / f"{prefix}.emapper.annotations"
|
|
seed_file = output_dir / f"{prefix}.emapper.seed_orthologs"
|
|
pfam_file = output_dir / f"{prefix}.emapper.pfam"
|
|
|
|
for f in [annot_file, seed_file, pfam_file]:
|
|
if f.exists():
|
|
self._record_file(f)
|
|
|
|
return {
|
|
"annotations": annot_file,
|
|
"seed_orthologs": seed_file,
|
|
"pfam": pfam_file,
|
|
}
|
|
|
|
def generate_report(self, output_dir: Path, stats: Dict) -> Path:
|
|
report_file = output_dir / "report.md"
|
|
params = self.config["parameters"]["emapper"]
|
|
|
|
report_content = f"""# 功能注释分析报告
|
|
|
|
## 基本信息
|
|
|
|
- **分析工具**: eggNOG-mapper v2.1.12
|
|
- **搜索模式**: {params.get('mode', 'diamond')}
|
|
- **E-value阈值**: {params.get('evalue', 0.001)}
|
|
- **分析时间**: {datetime.now().strftime("%Y-%m-%d %H:%M:%S")}
|
|
- **技能版本**: {self.config["skill"]["version"]}
|
|
|
|
## 统计结果
|
|
|
|
| 指标 | 数量 |
|
|
|------|------|
|
|
| 总基因数 | {stats.get('total_genes', 0)} |
|
|
| 有GO注释 | {stats.get('with_go', 0)} ({stats.get('with_go', 0)/stats.get('total_genes', 1)*100:.1f}%) |
|
|
| 有KEGG注释 | {stats.get('with_kegg', 0)} ({stats.get('with_kegg', 0)/stats.get('total_genes', 1)*100:.1f}%) |
|
|
| 有COG注释 | {stats.get('with_cog', 0)} ({stats.get('with_cog', 0)/stats.get('total_genes', 1)*100:.1f}%) |
|
|
| 有Pfam注释 | {stats.get('with_pfam', 0)} ({stats.get('with_pfam', 0)/stats.get('total_genes', 1)*100:.1f}%) |
|
|
| 独特GO terms | {stats.get('go_count', 0)} |
|
|
| 独特KEGG terms | {stats.get('kegg_count', 0)} |
|
|
| COG类别数 | {stats.get('cog_count', 0)} |
|
|
|
|
## 输出文件
|
|
|
|
| 文件 | 描述 |
|
|
|------|------|
|
|
| *.emapper.annotations | 主要注释结果 (TSV) |
|
|
| *.emapper.seed_orthologs | 种子直系同源 |
|
|
| *.emapper.pfam | Pfam结构域注释 |
|
|
|
|
## 可复现性
|
|
|
|
所有分析命令已记录在 `commands.sh` 中，
|
|
环境配置见 `environment.yml`。
|
|
|
|
## 校验信息
|
|
|
|
文件SHA256校验值见 `checksums.sha256`。
|
|
"""
|
|
|
|
with open(report_file, "w") as f:
|
|
f.write(report_content)
|
|
|
|
self._record_file(report_file)
|
|
return report_file
|
|
|
|
def write_reproducibility_bundle(self, output_dir: Path):
|
|
commands_file = output_dir / "commands.sh"
|
|
with open(commands_file, "w") as f:
|
|
f.write("#!/bin/bash\n# 功能注释分析复现命令\n\n")
|
|
for cmd in self.commands:
|
|
f.write(f"# {cmd['description']}\n")
|
|
f.write(f"{cmd['command']}\n\n")
|
|
os.chmod(commands_file, 0o755)
|
|
|
|
env_file = output_dir / "environment.yml"
|
|
env_config = {
|
|
"name": "annotation",
|
|
"channels": self.config["dependencies"]["conda"]["channels"],
|
|
"dependencies": self.config["dependencies"]["conda"]["packages"]
|
|
}
|
|
with open(env_file, "w") as f:
|
|
yaml.dump(env_config, f, default_flow_style=False)
|
|
|
|
checksums_file = output_dir / "checksums.sha256"
|
|
with open(checksums_file, "w") as f:
|
|
for filepath, info in self.checksums.items():
|
|
f.write(f"{info['sha256']} {filepath}\n")
|
|
|
|
def run(self, protein_fasta: Path, output_dir: Path, prefix: str = "annotation") -> Dict:
|
|
logger.info(f"Starting functional annotation with eggNOG-mapper")
|
|
|
|
self._record_file(protein_fasta, "input")
|
|
|
|
results = self.run_emapper(protein_fasta, output_dir, prefix)
|
|
|
|
stats = self._parse_annotations(results["annotations"])
|
|
self.stats = stats
|
|
|
|
report = self.generate_report(output_dir, stats)
|
|
self.write_reproducibility_bundle(output_dir)
|
|
|
|
logger.info(f"Annotation completed. Report: {report}")
|
|
|
|
return {
|
|
"status": "success",
|
|
"output_dir": str(output_dir),
|
|
"report": str(report),
|
|
"annotations": str(results["annotations"]),
|
|
"stats": stats,
|
|
}
|
|
|
|
def run_demo(self, output_dir: Path) -> Dict:
|
|
logger.info("Running in demo mode")
|
|
|
|
demo_dir = output_dir / "demo_data"
|
|
demo_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
protein_fasta = demo_dir / "demo_proteins.fa"
|
|
with open(protein_fasta, "w") as f:
|
|
f.write(">Gene001\nMKTLLLPLLLLVLLLLLVQGSISAQAAAAAA\n"
|
|
">Gene002\nMKTLLLPLLLLVLLLLLVQGSISAQBBBBBB\n"
|
|
">Gene003\nMKTLLLPLLLLVLLLLLVQGSISAQCCCCCC\n")
|
|
|
|
output_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
annot_file = output_dir / "annotation.emapper.annotations"
|
|
annot_file.write_text("""## emapper-version: 2.1.12
|
|
## command: emapper.py -i demo_proteins.fa --output annotation
|
|
#query seed_ortholog evalue score eggnog_OGs max_annot_lvl COG_category Description Preferred_name GOs EC KEGG_ko KEGG_Pathway KEGG_Module KEGG_Reaction KEGG_rclass BRITE KEGG_TC CAZy BiGG_Reaction PFAMs
|
|
Gene001 10090.ENSMUSP00000000001 1e-50 200.0 COG0001@1|root Eukaryota J Protein description GENE1 GO:0005524,GO:0004674 3.1.3.16 ko:K00001 path:ko00010 M00001 R00001 RC00001 ko00001 - GH1 - PF00001
|
|
Gene002 10090.ENSMUSP00000000002 2e-45 180.0 COG0002@1|root Eukaryota K Another protein GENE2 GO:0003677 - ko:K00002 path:ko00020 M00002 R00002 RC00002 ko00002 - - - PF00002
|
|
Gene003 10090.ENSMUSP00000000003 5e-40 160.0 COG0003@1|root Eukaryota L Third protein GENE3 - - - - - - - - - - - PF00003
|
|
""")
|
|
|
|
seed_file = output_dir / "annotation.emapper.seed_orthologs"
|
|
seed_file.write_text("Gene001\t10090.ENSMUSP00000000001\t1e-50\t200.0\n"
|
|
"Gene002\t10090.ENSMUSP00000000002\t2e-45\t180.0\n"
|
|
"Gene003\t10090.ENSMUSP00000000003\t5e-40\t160.0\n")
|
|
|
|
pfam_file = output_dir / "annotation.emapper.pfam"
|
|
pfam_file.write_text("Gene001\tPF00001\nGene002\tPF00002\nGene003\tPF00003\n")
|
|
|
|
stats = self._parse_annotations(annot_file)
|
|
report = self.generate_report(output_dir, stats)
|
|
self.write_reproducibility_bundle(output_dir)
|
|
|
|
logger.info(f"Demo annotation completed. Report: {report}")
|
|
|
|
return {
|
|
"status": "success",
|
|
"output_dir": str(output_dir),
|
|
"report": str(report),
|
|
"annotations": str(annot_file),
|
|
"stats": stats,
|
|
}
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description="Functional Annotation Skill - eggNOG-mapper")
|
|
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("--config", "-c", type=Path, help="配置文件路径")
|
|
parser.add_argument("--cpu", type=int, help="CPU线程数")
|
|
parser.add_argument("--mode", choices=["diamond", "mmseqs", "hmmer"], help="搜索模式")
|
|
parser.add_argument("--evalue", type=float, help="E-value阈值")
|
|
parser.add_argument("--demo", action="store_true", help="运行Demo模式")
|
|
|
|
args = parser.parse_args()
|
|
|
|
skill = AnnotationSkill(args.config)
|
|
|
|
if args.cpu:
|
|
skill.config["parameters"]["emapper"]["cpu"] = args.cpu
|
|
if args.mode:
|
|
skill.config["parameters"]["emapper"]["mode"] = args.mode
|
|
if args.evalue:
|
|
skill.config["parameters"]["emapper"]["evalue"] = args.evalue
|
|
|
|
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())
|
|
|