242 lines
8.1 KiB
Python
242 lines
8.1 KiB
Python
|
#!/usr/bin/env python3
|
||
|
|
"""
|
||
|
|
KaKs Analysis Utilities
|
||
|
|
Ka/Ks分析工具函数
|
||
|
|
"""
|
||
|
|
|
||
|
|
import hashlib
|
||
|
|
import logging
|
||
|
|
import re
|
||
|
|
import shutil
|
||
|
|
import subprocess
|
||
|
|
from pathlib import Path
|
||
|
|
from typing import Dict, List, Optional, Tuple
|
||
|
|
|
||
|
|
from Bio import SeqIO
|
||
|
|
|
||
|
|
logger = logging.getLogger(__name__)
|
||
|
|
|
||
|
|
|
||
|
|
def compute_checksum(filepath: Path) -> str:
|
||
|
|
"""计算文件SHA256校验值"""
|
||
|
|
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 check_tool(tool_name: str) -> Optional[str]:
|
||
|
|
"""检查工具是否可用"""
|
||
|
|
path = shutil.which(tool_name)
|
||
|
|
if path:
|
||
|
|
return path
|
||
|
|
return None
|
||
|
|
|
||
|
|
|
||
|
|
def get_tool_version(tool_name: str) -> str:
|
||
|
|
"""获取工具版本信息"""
|
||
|
|
try:
|
||
|
|
if tool_name == "muscle":
|
||
|
|
result = subprocess.run(
|
||
|
|
["muscle", "-version"],
|
||
|
|
capture_output=True, text=True, timeout=5
|
||
|
|
)
|
||
|
|
return result.stdout.strip().split('\n')[0] if result.stdout else "unknown"
|
||
|
|
elif tool_name == "yn00" or tool_name == "codeml":
|
||
|
|
result = subprocess.run(
|
||
|
|
[tool_name],
|
||
|
|
capture_output=True, text=True, timeout=5
|
||
|
|
)
|
||
|
|
output = result.stderr.strip() or result.stdout.strip()
|
||
|
|
lines = output.split('\n')
|
||
|
|
for line in lines:
|
||
|
|
if "version" in line.lower() or "paml" in line.lower():
|
||
|
|
return line.strip()
|
||
|
|
return "PAML (version unknown)"
|
||
|
|
elif tool_name == "pal2nal":
|
||
|
|
return "pal2nal v14.1"
|
||
|
|
except Exception as e:
|
||
|
|
logger.warning(f"Failed to get {tool_name} version: {e}")
|
||
|
|
return "not_found"
|
||
|
|
|
||
|
|
|
||
|
|
def parse_cds_fasta(fasta_path: Path) -> Dict[str, str]:
|
||
|
|
"""解析CDS FASTA文件，返回基因ID到序列的字典"""
|
||
|
|
sequences = {}
|
||
|
|
for record in SeqIO.parse(fasta_path, "fasta"):
|
||
|
|
gene_id = record.id
|
||
|
|
sequences[gene_id] = str(record.seq)
|
||
|
|
return sequences
|
||
|
|
|
||
|
|
|
||
|
|
def parse_ortholog_pairs(pairs_file: Path) -> List[Tuple[str, str]]:
|
||
|
|
"""解析同源基因对文件，格式: gene1\tgene2"""
|
||
|
|
pairs = []
|
||
|
|
with open(pairs_file, 'r') as f:
|
||
|
|
for line in f:
|
||
|
|
line = line.strip()
|
||
|
|
if not line or line.startswith('#'):
|
||
|
|
continue
|
||
|
|
parts = line.split('\t')
|
||
|
|
if len(parts) >= 2:
|
||
|
|
gene1 = parts[0].strip()
|
||
|
|
gene2 = parts[1].strip()
|
||
|
|
pairs.append((gene1, gene2))
|
||
|
|
return pairs
|
||
|
|
|
||
|
|
|
||
|
|
def extract_sequences(cds_dict: Dict[str, str], gene_ids: List[str], output_path: Path) -> int:
|
||
|
|
"""从CDS字典中提取指定基因的序列并写入文件"""
|
||
|
|
count = 0
|
||
|
|
with open(output_path, 'w') as f:
|
||
|
|
for gene_id in gene_ids:
|
||
|
|
if gene_id in cds_dict:
|
||
|
|
seq = cds_dict[gene_id]
|
||
|
|
f.write(f">{gene_id}\n")
|
||
|
|
for i in range(0, len(seq), 60):
|
||
|
|
f.write(seq[i:i+60] + '\n')
|
||
|
|
count += 1
|
||
|
|
else:
|
||
|
|
logger.warning(f"Gene {gene_id} not found in CDS file")
|
||
|
|
return count
|
||
|
|
|
||
|
|
|
||
|
|
def run_muscle_alignment(input_fasta: Path, output_aln: Path, maxiters: int = 2, cpu: int = 8) -> bool:
|
||
|
|
"""运行MUSCLE进行蛋白序列比对"""
|
||
|
|
muscle_path = check_tool("muscle")
|
||
|
|
if not muscle_path:
|
||
|
|
logger.error("muscle not found in PATH")
|
||
|
|
return False
|
||
|
|
|
||
|
|
try:
|
||
|
|
cmd = [
|
||
|
|
"muscle", "-align", str(input_fasta), "-output", str(output_aln),
|
||
|
|
"-maxiters", str(maxiters), "-threads", str(cpu)
|
||
|
|
]
|
||
|
|
logger.info(f"Running MUSCLE: {' '.join(cmd)}")
|
||
|
|
subprocess.run(cmd, check=True, capture_output=True, text=True)
|
||
|
|
return output_aln.exists() and output_aln.stat().st_size > 0
|
||
|
|
except Exception as e:
|
||
|
|
logger.error(f"MUSCLE error: {e}")
|
||
|
|
return False
|
||
|
|
|
||
|
|
|
||
|
|
def run_pal2nal(protein_aln: Path, cds_fasta: Path, output_codon: Path, output_format: str = "fasta") -> bool:
|
||
|
|
"""运行pal2nal将蛋白比对转换为密码子比对"""
|
||
|
|
pal2nal_path = check_tool("pal2nal.pl") or check_tool("pal2nal")
|
||
|
|
if not pal2nal_path:
|
||
|
|
logger.error("pal2nal.pl not found in PATH")
|
||
|
|
return False
|
||
|
|
|
||
|
|
try:
|
||
|
|
cmd = [
|
||
|
|
"perl", str(pal2nal_path), str(protein_aln), str(cds_fasta),
|
||
|
|
"-output", output_format, "-nogap", "-nomismatch"
|
||
|
|
]
|
||
|
|
logger.info(f"Running pal2nal: {' '.join(cmd)}")
|
||
|
|
with open(output_codon, 'w') as out_f:
|
||
|
|
subprocess.run(cmd, check=True, stdout=out_f, stderr=subprocess.PIPE, text=True)
|
||
|
|
return output_codon.exists() and output_codon.stat().st_size > 0
|
||
|
|
except Exception as e:
|
||
|
|
logger.error(f"pal2nal error: {e}")
|
||
|
|
return False
|
||
|
|
|
||
|
|
|
||
|
|
def create_yn00_ctl(seqfile: Path, outfile: Path, genetic_code: int = 0) -> Path:
|
||
|
|
"""创建yn00的控制文件"""
|
||
|
|
ctl_path = seqfile.parent / "yn00.ctl"
|
||
|
|
ctl_content = f"""seqfile = {seqfile}
|
||
|
|
outfile = {outfile}
|
||
|
|
noisy = 9
|
||
|
|
verbose = 1
|
||
|
|
icode = {genetic_code}
|
||
|
|
weighting = 0
|
||
|
|
commonf3x4 = 0
|
||
|
|
"""
|
||
|
|
with open(ctl_path, 'w') as f:
|
||
|
|
f.write(ctl_content)
|
||
|
|
return ctl_path
|
||
|
|
|
||
|
|
|
||
|
|
def run_yn00(codon_aln: Path, output_dir: Path, genetic_code: int = 0) -> Optional[Path]:
|
||
|
|
"""运行yn00计算Ka/Ks"""
|
||
|
|
yn00_path = check_tool("yn00")
|
||
|
|
if not yn00_path:
|
||
|
|
logger.error("yn00 not found in PATH")
|
||
|
|
return None
|
||
|
|
|
||
|
|
try:
|
||
|
|
output_file = output_dir / "yn00.out"
|
||
|
|
ctl_file = create_yn00_ctl(codon_aln, output_file, genetic_code)
|
||
|
|
cmd = [yn00_path, str(ctl_file)]
|
||
|
|
logger.info(f"Running yn00: {' '.join(cmd)}")
|
||
|
|
subprocess.run(cmd, cwd=output_dir, check=True, capture_output=True, text=True)
|
||
|
|
return output_file if output_file.exists() and output_file.stat().st_size > 0 else None
|
||
|
|
except Exception as e:
|
||
|
|
logger.error(f"yn00 error: {e}")
|
||
|
|
return None
|
||
|
|
|
||
|
|
|
||
|
|
def parse_yn00_output(yn00_file: Path) -> List[Dict]:
|
||
|
|
"""解析yn00输出文件，提取Ka/Ks结果"""
|
||
|
|
results = []
|
||
|
|
current_pair = {}
|
||
|
|
|
||
|
|
with open(yn00_file, 'r') as f:
|
||
|
|
content = f.read()
|
||
|
|
|
||
|
|
lines = content.split('\n')
|
||
|
|
i = 0
|
||
|
|
while i < len(lines):
|
||
|
|
line = lines[i].strip()
|
||
|
|
if line.startswith('(') and 'vs.' in line:
|
||
|
|
match = re.search(r'\(([^)]+)\s+vs\.\s+([^)]+)\)', line)
|
||
|
|
if match:
|
||
|
|
gene1 = match.group(1).strip()
|
||
|
|
gene2 = match.group(2).strip()
|
||
|
|
current_pair = {'gene1': gene1, 'gene2': gene2}
|
||
|
|
|
||
|
|
j = i + 1
|
||
|
|
while j < len(lines) and j < i + 20:
|
||
|
|
next_line = lines[j].strip()
|
||
|
|
if 'Ka' in next_line and 'Ks' in next_line:
|
||
|
|
if j + 1 < len(lines):
|
||
|
|
values_line = lines[j + 1].strip()
|
||
|
|
parts = values_line.split()
|
||
|
|
if len(parts) >= 4:
|
||
|
|
try:
|
||
|
|
current_pair['Ka'] = float(parts[0])
|
||
|
|
current_pair['Ks'] = float(parts[1])
|
||
|
|
current_pair['Ka_Ks'] = float(parts[2]) if parts[2] != 'nan' else float('nan')
|
||
|
|
current_pair['method'] = parts[3] if len(parts) > 3 else 'YN00'
|
||
|
|
results.append(current_pair.copy())
|
||
|
|
except ValueError:
|
||
|
|
pass
|
||
|
|
break
|
||
|
|
j += 1
|
||
|
|
i += 1
|
||
|
|
|
||
|
|
return results
|
||
|
|
|
||
|
|
|
||
|
|
def calculate_divergence_time(ks_values: List[float], neutral_rate: float = 7.08e-9) -> Dict:
|
||
|
|
"""计算分歧时间"""
|
||
|
|
valid_ks = [ks for ks in ks_values if ks > 0 and not (ks != ks)]
|
||
|
|
|
||
|
|
if not valid_ks:
|
||
|
|
return {'mean_ks': None, 'median_ks': None, 'divergence_time_my': None}
|
||
|
|
|
||
|
|
mean_ks = sum(valid_ks) / len(valid_ks)
|
||
|
|
median_ks = sorted(valid_ks)[len(valid_ks) // 2]
|
||
|
|
|
||
|
|
# T = Ks / (2 * neutral_rate)
|
||
|
|
divergence_time = mean_ks / (2 * neutral_rate) / 1e6 # 转换为百万年
|
||
|
|
|
||
|
|
return {
|
||
|
|
'mean_ks': mean_ks,
|
||
|
|
'median_ks': median_ks,
|
||
|
|
'divergence_time_my': divergence_time,
|
||
|
|
'pair_count': len(valid_ks)
|
||
|
|
}
|